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

mimicpy/force_matching/bonded_forces.py

@@ -0,0 +1,628 @@
+import numpy as np
+
+
+def compute_bond_force(xi, xj, bond_length, force_constant, 
+                       flag='force', derivative=None):
+    """
+    Computes forces or derivatives related to a bond.
+    For harmonic bonds: V(r) = 1/2 * k * (r - r_0)^2
+    
+    Args:
+        xi, xj: Coordinates of the two atoms
+        bond_length: Equilibrium bond length (r_0)
+        force_constant: Force constant (k)
+        flag: 'force' or 'derivative'
+        derivative: 'force_constant' or 'bond_length'
+    
+    Returns:
+        tuple: (fi, fj) Computed forces or derivatives
+    """
+    # Compute distance vector and length
+    R_ij = xj - xi
+    r_ij = get_magnitude(R_ij)
+    
+    # Check for zero length vector
+    if r_ij == 0:
+        raise ValueError("Zero length bond vector encountered")
+    
+    # Compute unit vector
+    R_ij_unit = R_ij / r_ij
+    
+    # Compute force derivatives with respect to atomic positions
+    # These are the derivatives of r with respect to atomic positions
+    dr_dri = -R_ij_unit  # Derivative of r with respect to ri
+    dr_drj = R_ij_unit  # Derivative of r with respect to rj
+    
+    if flag == 'force':
+        # Force is negative gradient of potential: F = -k * (r - r_0) * dr/dr
+        F_i = -force_constant * (r_ij - bond_length) * dr_dri
+        F_j = -force_constant * (r_ij - bond_length) * dr_drj
+        return F_i, F_j
+    
+    elif flag == 'derivative':
+        if derivative == 'force_constant':
+            # Derivative of force with respect to force constant
+            # dF/dk = -(r - r_0) * dr/dr
+            dk_i = -(r_ij - bond_length) * dr_dri
+            dk_j = -(r_ij - bond_length) * dr_drj
+            return dk_i, dk_j
+            
+        elif derivative == 'bond_length':
+            # Derivative of force with respect to equilibrium bond length
+            # dF/dr_0 = k * dr/dr
+            dr0_i = force_constant * dr_dri
+            dr0_j = force_constant * dr_drj
+            return dr0_i, dr0_j
+
+def compute_angle_force(xi, xj, xk, angle, force_constant, flag='force', derivative=None):
+    """
+    Computes forces or derivatives related to an angle.
+    For harmonic angles: V(theta) = 1/2 * k * (theta - theta_0)^2
+    
+    Args:
+        xi, xj, xk: Coordinates of the three atoms
+        angle: Equilibrium angle (theta_0)
+        force_constant: Force constant (k)
+        flag: 'force' or 'derivative'
+        derivative: 'force_constant' or 'angle'
+    
+    Returns:
+        tuple: (fi, fj, fk) Computed forces or derivatives
+    """
+    # Compute distance vectors and distances
+    R_ij = xj - xi
+    R_kj = xj - xk
+    r_ij = np.linalg.norm(R_ij)
+    r_kj = np.linalg.norm(R_kj)
+    
+    # Check for zero length vectors
+    if r_ij == 0 or r_kj == 0:
+        raise ValueError("Zero length vector encountered in angle force computation")
+    
+    # Compute angle
+    dot_r_ij_kj = np.dot(R_ij, R_kj) 
+    prod_r_ij_kj = r_ij * r_kj 
+    cos_theta = np.clip(dot_r_ij_kj/prod_r_ij_kj, -1.0, 1.0)
+    theta = np.arccos(cos_theta)
+    
+    # Compute sin(theta) safely
+    sin_theta = np.sqrt(1.0 - cos_theta**2)
+    if sin_theta < 1e-10:  # Near linear angles
+        raise ValueError("Near linear angle encountered")
+    
+    # Compute force derivatives with respect to atomic positions
+    # These are the derivatives of theta with respect to atomic positions
+    f_i = 1/(r_ij * sin_theta) * (R_kj/r_kj - cos_theta * R_ij/r_ij)
+    f_k = 1/(r_kj * sin_theta) * (R_ij/r_ij - cos_theta * R_kj/r_kj)
+    f_j = -f_i - f_k  # Force on central atom
+    
+    if flag == 'force':
+        # Force is negative gradient of potential: F = -k * (theta - theta_0) * dtheta/dr
+        F_i = -force_constant * (theta - angle) * f_i
+        F_k = -force_constant * (theta - angle) * f_k
+        F_j = -force_constant * (theta - angle) * f_j
+        return F_i, F_j, F_k
+    
+    elif flag == 'derivative':
+        if derivative == 'force_constant':
+            # Derivative of force with respect to force constant
+            # dF/dk = -(theta - theta_0) * dtheta/dr
+            dk_i = -(theta - angle) * f_i
+            dk_k = -(theta - angle) * f_k
+            dk_j = -(theta - angle) * f_j
+            return dk_i, dk_j, dk_k
+            
+        elif derivative == 'angle':
+            # Derivative of force with respect to equilibrium angle
+            # dF/dtheta_0 = k * dtheta/dr
+            dtheta_i = force_constant * f_i
+            dtheta_k = force_constant * f_k
+            dtheta_j = force_constant * f_j
+            return dtheta_i, dtheta_j, dtheta_k
+        
+
+def get_magnitude(vector):
+    return np.linalg.norm(vector)
+
+def distance_vector(R1, R2):
+    r = R1 - R2
+    d = get_magnitude(r)
+    return d, r
+
+def compute_RB_dihedral_force(xi, xj, xk, xl, 
+                              force_constant, flag='force', derivative=None):
+    """
+    Computes forces or derivatives related to a Ryckaert-Bellemans (RB) dihedral angle.
+    
+    The RB potential is defined as:
+    V(φ) = C0 + C1*cos(φ-180°) + C2*cos²(φ-180°) + C3*cos³(φ-180°) + C4*cos⁴(φ-180°) + C5*cos⁵(φ-180°)
+    
+    where φ is the dihedral angle and C0-C5 are the force constants.
+    
+    Args:
+        xi, xj, xk, xl: Coordinates of the four atoms defining the dihedral angle
+        force_constant: Array of force constants [C0, C1, C2, C3, C4, C5]
+        flag: 'force' or 'derivative'
+        derivative: 'C0', 'C1', 'C2', 'C3', 'C4', or 'C5' for parameter derivatives
+    
+    Returns:
+        tuple: (fi, fj, fk, fl) Computed forces or derivatives for each atom
+    """
+
+    # Compute distance vectors and distances
+    R_ij = xj - xi
+    R_kj = xj - xk
+    R_kl = xl - xk
+    
+    r_kj = np.linalg.norm(R_kj)
+
+    # Projection of vectors on the plane orthogonal to R_kj
+    R_im = R_ij - np.dot(R_ij, R_kj) * R_kj / r_kj**2
+    R_ln = -R_kl + np.dot(R_kl, R_kj) * R_kj / r_kj**2
+
+    # Magnitudes of the projections
+    r_im = np.linalg.norm(R_im)
+    r_ln = np.linalg.norm(R_ln)
+
+    # Compute cross product and sign of the dihedral angle
+    cosphi = np.dot(R_im, R_ln) / (r_im * r_ln)
+    cosphi = np.clip(cosphi, -1.0, 1.0)  # Ensure cosphi is in valid range
+    
+    # For RB potential, we need cos(phi-180°) = -cos(phi)
+    cosphi_rb = -cosphi
+    
+    # Calculate forces
+    di = (R_ln/r_ln - R_im/r_im * cosphi) * 1/r_im
+    dl = (R_im/r_im - R_ln/r_ln * cosphi) * 1/r_ln
+    dj = (np.dot(R_ij, R_kj) / r_kj**2 - 1.0) * di - (np.dot(R_kl, R_kj) / r_kj**2) * dl
+    dk = -di - dj - dl
+
+    if flag == 'force':
+        # RB potential: V(phi) = C0 + C1*cos(phi-180°) + C2*cos²(phi-180°) + C3*cos³(phi-180°) + C4*cos⁴(phi-180°) + C5*cos⁵(phi-180°)
+        # Force factor is the negative derivative of V with respect to cos(phi-180°)
+        factor = -(force_constant[1]  # C1
+                  + force_constant[2] * 2 * cosphi_rb  # C2
+                  + force_constant[3] * 3 * cosphi_rb**2  # C3
+                  + force_constant[4] * 4 * cosphi_rb**3  # C4
+                  + force_constant[5] * 5 * cosphi_rb**4)  # C5
+        
+        fi = factor * di
+        fj = factor * dj
+        fk = factor * dk
+        fl = factor * dl
+        return fi, fj, fk, fl
+    
+    elif flag=='derivative':
+        # For derivatives, we need dF/dC_n where F is the force
+        # F = -d/d(cos(phi-180°)) of V * di
+        # So dF/dC_n = -d/d(cos(phi-180°)) of (C_n * cos^n(phi-180°)) * di
+        if derivative == 'C0':
+            factor = 0.0  # C0 doesn't contribute to forces
+        elif derivative == 'C1':
+            factor = -1.0  # -d/d(cos(phi-180°)) of C1*cos(phi-180°)
+        elif derivative == 'C2':
+            factor = -2 * cosphi_rb  # -d/d(cos(phi-180°)) of C2*cos²(phi-180°)
+        elif derivative == 'C3':
+            factor = -3 * cosphi_rb**2  # -d/d(cos(phi-180°)) of C3*cos³(phi-180°)
+        elif derivative == 'C4':
+            factor = -4 * cosphi_rb**3  # -d/d(cos(phi-180°)) of C4*cos⁴(phi-180°)
+        elif derivative == 'C5':
+            factor = -5 * cosphi_rb**4  # -d/d(cos(phi-180°)) of C5*cos⁵(phi-180°)
+        else:
+            raise ValueError(f'Unknown derivative parameter: {derivative}')
+        
+        # The derivatives of the forces with respect to the coefficients
+        fi = factor * di
+        fj = factor * dj
+        fk = factor * dk
+        fl = factor * dl
+        return fi, fj, fk, fl
+    
+def compute_dihedral_force(xi, xj, xk, xl, phase, force_constant, m, flag='force', derivative=None):
+    """
+    Computes forces or derivatives related to a dihedral angle.
+    For periodic dihedrals (types 1, 4, 9): V(phi) = k[1 + cos(n*phi - phi_0)]
+   
+    Args:
+        xi, xj, xk, xl: Coordinates of the four atoms
+        phase: Phase shift (phi_0)
+        force_constant: Force constant (k)
+        m: Multiplicity (n)
+        flag: 'force' or 'derivative'
+        derivative: Not used for periodic dihedrals
+    
+    Returns:
+        tuple: (fi, fj, fk, fl) Computed forces or derivatives
+    """
+    # Compute distance vectors and distances
+    R_ij = xj - xi
+    R_kj = xj - xk
+    R_kl = xl - xk
+    
+    r_kj = np.linalg.norm(R_kj)
+    
+    R_mj = np.cross(R_ij, R_kj)
+    R_nk = np.cross(R_kj, R_kl)
+    r_mj = np.linalg.norm(R_mj)
+    r_nk = np.linalg.norm(R_nk)
+    
+    # Projection of vectors on the plane orthogonal to R_kj
+    R_im = R_ij - np.dot(R_ij, R_kj) * R_kj / r_kj**2
+    R_ln = -R_kl + np.dot(R_kl, R_kj) * R_kj / r_kj**2
+
+    # Magnitudes of the projections
+    r_im = np.linalg.norm(R_im)
+    r_ln = np.linalg.norm(R_ln)
+
+    # Compute cross product and sign of the dihedral angle
+    cosphi = np.dot(R_im, R_ln) / (r_im * r_ln)
+    cosphi = np.clip(cosphi, -1.0, 1.0)  # Ensure cosphi is in valid range
+    phi = np.sign(np.dot(R_ij, R_nk))*np.arccos(cosphi)
+    sinphi = np.sin(m*phi - phase)
+    
+    # Calculate force derivatives with respect to atomic positions
+    di = m * sinphi * r_kj * R_mj / r_mj**2 
+    dl = - m * sinphi * r_kj * R_nk / r_nk**2
+    dj = (np.dot(R_ij, R_kj) / r_kj**2 - 1.0) * di - (np.dot(R_kl, R_kj) / r_kj**2) * dl
+    dk = -di - dj - dl
+
+    if flag == 'force':
+        # Force is negative gradient of potential
+        F_i = -force_constant * di
+        F_l = -force_constant * dl
+        F_j = -force_constant * dj
+        F_k = -force_constant * dk
+        return F_i, F_j, F_k, F_l
+    
+    elif flag=='derivative':
+        # For periodic dihedrals, we only need derivative with respect to force constant
+        # dF/dk = -di, -dj, -dk, -dl
+        return -di, -dj, -dk, -dl
+
+def get_jac_indx(atom_idx):
+        return  3*atom_idx
+
+def compute_bonds_forces(bonds, qm_coordinates, 
+                         ff_optimize,
+                          bond2params, flag='force'):
+    
+    function_map = {
+        1: compute_bond_force,
+    }
+    if flag == 'force':
+        forces = np.zeros(qm_coordinates.shape)
+        for bond in bonds:
+            idxs = bond2params.get(bond['index'])
+            b = bond['parameters'][0]
+            kb = bond['parameters'][1]
+            if idxs[0] is not None:
+                b = ff_optimize[idxs[0]]
+            if idxs[1] is not None:
+                kb = ff_optimize[idxs[1]]
+            
+            # Get atom coordinates using global indices
+            i_coords = qm_coordinates[bond['atoms'][0]]
+            j_coords = qm_coordinates[bond['atoms'][1]]
+            
+            fi, fj = function_map.get(bond['function'])(i_coords, j_coords,
+                                                      b, kb,
+                                                      flag=flag)
+            
+            # Add forces to the correct atoms
+            forces[bond['atoms'][0]] += fi
+            forces[bond['atoms'][1]] += fj
+
+        return forces
+    elif flag == 'derivative':
+        jac_mat = np.zeros((3*qm_coordinates.shape[0],ff_optimize.shape[0]))
+
+        for bond in bonds:
+            if not bond['optimize']:
+                continue
+                
+            idxs = bond2params.get(bond['index'])
+            b = bond['parameters'][0]
+            kb = bond['parameters'][1]
+
+            if idxs[0] is not None:
+                b = ff_optimize[idxs[0]]
+            if idxs[1] is not None:
+                kb = ff_optimize[idxs[1]]
+
+            idx0 = get_jac_indx(bond['atoms'][0])
+            idx1 = get_jac_indx(bond['atoms'][1])
+
+            if idxs[0] is not None:
+                bi, bj = function_map.get(bond['function'])(qm_coordinates[bond['atoms'][0]], 
+                                                          qm_coordinates[bond['atoms'][1]],
+                                                          b, kb,
+                                                          flag=flag, derivative='bond_length')
+                
+                jac_mat[idx0:idx0+3, idxs[0]] += bi
+                jac_mat[idx1:idx1+3, idxs[0]] += bj
+            
+            if idxs[1] is not None:
+                ki, kj = function_map.get(bond['function'])(qm_coordinates[bond['atoms'][0]], 
+                                                          qm_coordinates[bond['atoms'][1]],
+                                                          b, kb,
+                                                          flag=flag, derivative='force_constant')
+                
+                jac_mat[idx0:idx0+3, idxs[1]] += ki
+                jac_mat[idx1:idx1+3, idxs[1]] += kj
+
+        return jac_mat
+            
+
+
+def compute_angles_forces(angles, qm_coordinates, 
+                         ff_optimize,
+                          bond2params, flag='force'):
+    
+    function_map = {
+        1: compute_angle_force,
+    }
+    if flag == 'force':
+
+        forces = np.zeros(qm_coordinates.shape)
+
+        for angle in angles:
+            idxs = bond2params.get(angle['index'])
+            theta = angle['parameters'][0]
+            cth = angle['parameters'][1]
+            if idxs[0] is not None:
+                theta = ff_optimize[idxs[0]]
+            if idxs[1] is not None:
+                cth = ff_optimize[idxs[1]]
+            fi, fj, fk = function_map.get(angle['function'])(qm_coordinates[angle['atoms'][0]], 
+                                            qm_coordinates[angle['atoms'][1]], 
+                                            qm_coordinates[angle['atoms'][2]],
+                                            theta, cth,
+                                            flag=flag)
+            
+            forces[angle['atoms'][0]] += fi
+            forces[angle['atoms'][1]] += fj
+            forces[angle['atoms'][2]] += fk
+
+        return forces
+
+    elif flag == 'derivative':
+        jac_mat = np.zeros((3*qm_coordinates.shape[0],ff_optimize.shape[0]))
+
+        for angle in angles:
+
+            if not angle['optimize']:
+                continue
+            idxs = bond2params.get(angle['index'])
+            theta = angle['parameters'][0]
+            cth = angle['parameters'][1]
+            
+            if idxs[0] is not None:
+                theta = ff_optimize[idxs[0]]
+            if idxs[1] is not None:
+                cth = ff_optimize[idxs[1]]
+            
+            idx0 =  get_jac_indx(angle['atoms'][0])
+            idx1 = get_jac_indx(angle['atoms'][1])
+            idx2 = get_jac_indx(angle['atoms'][2])
+            
+            if idxs[0] is not None:
+                agi, agj, agk   = function_map.get(angle['function'])(qm_coordinates[angle['atoms'][0]], 
+                                    qm_coordinates[angle['atoms'][1]],
+                                    qm_coordinates[angle['atoms'][2]],
+                                        theta, cth,
+                                        flag='derivative',
+                                        derivative='angle')
+            
+                
+                jac_mat[idx0: idx0 +3, idxs[0]] += agi
+                jac_mat[idx1 : idx1 +3, idxs[0]] += agj
+                jac_mat[idx2 : idx2 +3, idxs[0]] += agk
+            
+            if idxs[1] is not None:
+                ki, kj, kk = function_map.get(angle['function'])(qm_coordinates[angle['atoms'][0]], 
+                                        qm_coordinates[angle['atoms'][1]],
+                                        qm_coordinates[angle['atoms'][2]],
+                                             theta, cth,
+                                            flag='derivative',
+                                            derivative='force_constant')
+                
+                jac_mat[idx0: idx0 +3, idxs[1]] += ki
+                jac_mat[idx1 : idx1 +3, idxs[1]] += kj
+                jac_mat[idx2 : idx2 +3, idxs[1]] += kk
+
+        return jac_mat
+    
+
+
+def compute_dihedrals_forces(dihedrals, qm_coordinates, 
+                            ff_optimize,
+                            bond2params, flag='force'):
+    
+    function_map = {
+        4: compute_dihedral_force,
+        3: compute_RB_dihedral_force
+    }  
+    
+
+    if flag == 'force':
+        forces = np.zeros(qm_coordinates.shape)
+        for dihedral in dihedrals:
+            if dihedral['function'] in [1, 4, 9]:
+                kd = dihedral['parameters'][1]
+                idxs = bond2params.get(dihedral['index'])
+                if idxs[1] is not None:
+                    kd = ff_optimize[idxs[1]]
+                fi, fj, fk, fl= compute_dihedral_force(qm_coordinates[dihedral['atoms'][0]], 
+                                                    qm_coordinates[dihedral['atoms'][1]], 
+                                                    qm_coordinates[dihedral['atoms'][2]], 
+                                                    qm_coordinates[dihedral['atoms'][3]], 
+                                                    dihedral['parameters'][0], kd,
+                                                        dihedral['parameters'][2], flag=flag)
+                forces[dihedral['atoms'][0]] += fi
+                forces[dihedral['atoms'][1]] += fj
+                forces[dihedral['atoms'][2]] += fk
+                forces[dihedral['atoms'][3]] += fl
+
+            elif dihedral['function'] == 3:
+                C_params = []
+                idxs = bond2params.get(dihedral['index'])
+                for i, c in enumerate(idxs):
+                    if c is not None:
+                        C_params.append(ff_optimize[c])
+                    else:
+                        C_params.append(dihedral['parameters'][i])
+                fi, fj, fk, fl= compute_RB_dihedral_force(qm_coordinates[dihedral['atoms'][0]], 
+                                                    qm_coordinates[dihedral['atoms'][1]], 
+                                                    qm_coordinates[dihedral['atoms'][2]], 
+                                                    qm_coordinates[dihedral['atoms'][3]], 
+                                                    C_params,flag=flag)
+            
+                forces[dihedral['atoms'][0]] += fi
+                forces[dihedral['atoms'][1]] += fj
+                forces[dihedral['atoms'][2]] += fk
+                forces[dihedral['atoms'][3]] += fl
+
+        return forces
+    
+    elif flag == 'derivative':
+        jac_mat = np.zeros((3*qm_coordinates.shape[0],ff_optimize.shape[0]))
+        for dihedral in dihedrals:
+            if not dihedral['optimize']:
+                    continue
+            idxs = bond2params.get(dihedral['index'])
+            
+            if dihedral['function'] in [1, 4, 9]:
+                # Only compute derivatives if the force constant is being optimized
+                if idxs[1] is not None:
+                    ki, kj, kk, kl= compute_dihedral_force(qm_coordinates[dihedral['atoms'][0]], 
+                                                        qm_coordinates[dihedral['atoms'][1]], 
+                                                        qm_coordinates[dihedral['atoms'][2]], 
+                                                        qm_coordinates[dihedral['atoms'][3]], 
+                                                        dihedral['parameters'][0],  ff_optimize[idxs[1]],
+                                                            dihedral['parameters'][2], flag='derivative')
+                    idx0 =  get_jac_indx(dihedral['atoms'][0])
+                    idx1 = get_jac_indx(dihedral['atoms'][1])
+                    idx2 = get_jac_indx(dihedral['atoms'][2])
+                    idx3 = get_jac_indx(dihedral['atoms'][3])
+                    
+                    jac_mat[idx0: idx0 +3, idxs[1]] += ki
+                    jac_mat[idx1 : idx1 +3, idxs[1]] += kj
+                    jac_mat[idx2 : idx2 +3, idxs[1]] += kk
+                    jac_mat[idx3 : idx3 +3, idxs[1]] += kl
+
+            elif dihedral['function'] == 3:
+                if not dihedral['optimize']:
+                    continue
+                C_params = []
+                idxs = bond2params.get(dihedral['index'])
+                for i, c in enumerate(idxs):
+                    if c is not None:
+                        C_params.append(ff_optimize[c])
+                    else:
+                        C_params.append(dihedral['parameters'][i])
+                
+                for i, c in enumerate(idxs):
+                    if c is None:
+                        continue
+                    ki, kj, kk, kl= compute_RB_dihedral_force(qm_coordinates[dihedral['atoms'][0]], 
+                                                        qm_coordinates[dihedral['atoms'][1]], 
+                                                        qm_coordinates[dihedral['atoms'][2]], 
+                                                        qm_coordinates[dihedral['atoms'][3]], 
+                                                        C_params, flag='derivative',derivative=f'C{i}')
+                    
+                    idx0 =  get_jac_indx(dihedral['atoms'][0])
+                    idx1 = get_jac_indx(dihedral['atoms'][1])
+                    idx2 = get_jac_indx(dihedral['atoms'][2])
+                    idx3 = get_jac_indx(dihedral['atoms'][3])
+                    jac_mat[idx0: idx0 +3, c] += ki
+                    jac_mat[idx1 : idx1 +3, c] += kj
+                    jac_mat[idx2 : idx2 +3, c] += kk
+                    jac_mat[idx3 : idx3 +3, c] += kl
+
+        return jac_mat
+
+def compute_bonded_forces(ff_optimize,
+                            qm_coordinates, 
+                            bonds,
+                            angles,
+                            dihedrals,
+                            bond2params,
+                            qm_atoms_count=None):
+    """
+    Compute bonded forces for all interactions in the QM region.
+    
+    Args:
+        ff_optimize (numpy.ndarray): Array of optimized force field parameters
+        qm_coordinates (numpy.ndarray): Array of coordinates (QM only or QM+MM)
+        bonds (list): List of bond dictionaries
+        angles (list): List of angle dictionaries
+        dihedrals (list): List of dihedral dictionaries
+        bond2params (dict): Mapping of interaction indices to parameter indices
+        qm_atoms_count (int, optional): Number of QM atoms if coordinates include MM atoms
+        
+    Returns:
+        numpy.ndarray: Array of forces for QM atoms only
+    """
+    # Initialize forces array for all atoms in coordinates
+    forces = np.zeros(qm_coordinates.shape)
+    
+    # Process bonds
+    b_forces = compute_bonds_forces(bonds, qm_coordinates, 
+                                  ff_optimize,
+                                  bond2params, 'force')
+    forces += b_forces
+    
+    # Process angles
+    a_forces = compute_angles_forces(angles, qm_coordinates, 
+                                   ff_optimize, 
+                                   bond2params, 'force')
+    forces += a_forces
+    
+    # Process dihedrals
+    d_forces = compute_dihedrals_forces(dihedrals, qm_coordinates, 
+                                      ff_optimize,
+                                      bond2params, 'force')
+    forces += d_forces
+    
+    # If qm_atoms_count is provided, return only QM atom forces
+    if qm_atoms_count is not None:
+        return forces[:qm_atoms_count]
+    
+    return forces
+
+
+def jacobian_ff(ff_optimize,
+                qm_coordinates,
+                bond2params,
+                bonds, angles, dihedrals,
+                qm_atoms_count=None):
+    
+    """
+    ff_optimize: b1, k1, b2, k2, ...
+    """
+
+    number_of_parameters = ff_optimize.shape[0]
+    # If qm_atoms_count is provided, compute Jacobian only for QM atoms
+    if qm_atoms_count is not None:
+        number_of_residuals = 3 * qm_atoms_count
+    else:
+        number_of_residuals = 3 * qm_coordinates.shape[0]
+    
+    # compute the bonded forces
+    jacobian_matrix = np.zeros((number_of_residuals,number_of_parameters))
+    
+    # Compute Jacobian for all atoms in coordinates
+    full_jacobian = np.zeros((3 * qm_coordinates.shape[0], number_of_parameters))
+    
+    full_jacobian += compute_bonds_forces(bonds, qm_coordinates,
+                                            ff_optimize, bond2params, flag='derivative')
+    full_jacobian += compute_angles_forces(angles, qm_coordinates,
+                                             ff_optimize, bond2params, flag='derivative')
+    full_jacobian += compute_dihedrals_forces(dihedrals, qm_coordinates,
+                                                ff_optimize, bond2params, flag='derivative')
+    
+    # If qm_atoms_count is provided, return only QM atom Jacobian
+    if qm_atoms_count is not None:
+        return full_jacobian[:3*qm_atoms_count, :]
+    
+    return full_jacobian