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

mimicpy/force_matching/dresp.py

@@ -0,0 +1,435 @@
+import numpy as np
+import multiprocessing as mp
+from ..scripts.fmdata import FMDataset
+from .qm_region import QMRegion
+
+
+def get_configurations(fmdata: FMDataset, begin: int, end: int, step: int, qm_region: 'QMRegion'):
+    """Get a range of configurations for DRESP, with QM data reordered to GROMACS topology order.
+
+    Args:
+        fmdata (FMDataset): The force matching dataset.
+        begin (int): Starting frame index.
+        end (int): Ending frame index (exclusive).
+        step (int): Step size for frame iteration.
+        qm_region (QMRegion): The QMRegion object providing GROMACS/CPMD mappings and QM atom definitions.
+
+    Returns:
+        list: A list of configuration dictionaries. Each dictionary contains:
+            'qm_coordinates': QM atom coordinates, ordered by GROMACS QM atom topology.
+            'mm_coordinates': MM atom coordinates, as fetched from FMDataset (GROMACS ID ordered).
+            'electric_potential': Electric potential at MM atom/point locations, as fetched.
+            'electric_field': Electric field at MM atom/point locations, as fetched.
+            'reference_charge': QM atom Hirshfeld charges, ordered by GROMACS QM atom topology.
+            'qm_cpmd_ids_fmdata_order': Original CPMD IDs for QM atoms from FMDataset for this frame.
+            'mm_gmx_ids_fmdata_order': Original GROMACS IDs for MM atoms/points from FMDataset for this frame.
+    """
+    configurations = []
+    
+    # Target order for QM properties: GROMACS QM atom indices (1-based)
+    target_gmx_qm_indices = qm_region.qm_atoms.index # These are 1-based GROMACS indices
+    gmx_to_cpmd_map = qm_region.gmx_to_cpmd_map # Maps 1-based GMX to 1-based CPMD
+
+    # We need to build the reordering map once, assuming atom IDs are consistent across frames in fmdata.
+    # Get the CPMD IDs for QM atoms from the first frame in fmdata to establish the source order.
+    cpmd_ids_fmdata_order_qm = fmdata.get_configuration_properties(begin, 'id', 'qm')
+
+    qm_reorder_indices = np.zeros(len(target_gmx_qm_indices), dtype=int)
+    for i, gmx_idx_1_based in enumerate(target_gmx_qm_indices):
+        target_cpmd_id_1_based = gmx_to_cpmd_map.get(gmx_idx_1_based)
+        if target_cpmd_id_1_based is None:
+            raise ValueError(f"GROMACS QM atom {gmx_idx_1_based} not found in gmx_to_cpmd_map.")
+        
+        fm_idx_arr = np.where(cpmd_ids_fmdata_order_qm == target_cpmd_id_1_based)[0]
+        if len(fm_idx_arr) > 0:
+            qm_reorder_indices[i] = fm_idx_arr[0]
+        else:
+            raise ValueError(
+                f"CPMD ID {target_cpmd_id_1_based} (for GMX QM atom {gmx_idx_1_based}) "
+                f"not found in FMDataset's list of QM CPMD IDs: {cpmd_ids_fmdata_order_qm}"
+            )
+
+    for idx in range(begin, end, step):
+        config = dict()
+        
+        # QM data (ordered by CPMD IDs in fmdata)
+        raw_qm_coords = fmdata.get_configuration_properties(idx, 'coordinate', 'qm')
+        raw_qm_ref_charges = fmdata.get_configuration_properties(idx, 'hirshfeld_charge', 'qm')
+        # Store original QM CPMD IDs from fmdata for this frame if needed for debugging/verification
+        config['qm_cpmd_ids_fmdata_order'] = fmdata.get_configuration_properties(idx, 'id', 'qm')
+
+        # Reorder QM data to match GROMACS QM atom topology order
+        config['qm_coordinates'] = raw_qm_coords[qm_reorder_indices]
+        config['reference_charge'] = raw_qm_ref_charges[qm_reorder_indices]
+        
+        # MM data (ordered by GROMACS IDs in fmdata for MM region)
+        # We assume this order is what DRESP expects or can work with directly.
+        config['mm_coordinates'] = fmdata.get_configuration_properties(idx, 'coordinate', 'mm')
+        config['electric_potential'] = fmdata.get_configuration_properties(idx, 'electric_potential', 'mm').squeeze()
+        config['electric_field'] = fmdata.get_configuration_properties(idx, 'electric_field', 'mm')
+        # Store original MM GROMACS IDs from fmdata for this frame if needed
+        config['mm_gmx_ids_fmdata_order'] = fmdata.get_configuration_properties(idx, 'id', 'mm')
+        
+        configurations.append(config)
+    return configurations
+
+def compute_potential_set_charges(charges, charge_positions, x2):
+    R_ij = x2 - charge_positions 
+    r_ij = np.linalg.norm(R_ij, axis=1)
+    potential = charges * (1/r_ij)
+    return np.sum(potential)
+    
+
+def compute_electric_field_set_charges(charges, charge_positions, x2):
+    R_ij = x2 - charge_positions 
+    r_ij = np.linalg.norm(R_ij, axis=1)
+    dem = np.reshape(r_ij**3, (-1, 1))
+    charges = charges.reshape(-1,1)
+    electric_field = R_ij * (charges) * (1/dem)
+    return np.sum(electric_field, axis=0)
+
+def compute_diff_electric_field(charges, charge_positions, 
+                       sr_positions, sr_electric_field):
+    
+    Emm = np.array(sr_positions)
+    for i, sr_position in enumerate(sr_positions):
+        Emm[i] = compute_electric_field_set_charges(charges, charge_positions, sr_position)
+
+    Eroh = np.sum((sr_electric_field)**2)
+    E_diff = np.sum((Emm - sr_electric_field)**2)
+    return Eroh, E_diff 
+
+def compute_diff_potential(charges, charge_positions, 
+                       sr_positions, sr_potential):
+    Vmm = np.zeros(sr_positions.shape[0])
+    for i, sr_position in enumerate(sr_positions):
+        Vmm[i] = compute_potential_set_charges(charges, charge_positions, sr_position)
+
+    V_diff = np.sum((Vmm - sr_potential)**2) 
+    Vroh = np.sum((sr_potential)**2)
+    
+  
+    return Vroh, V_diff 
+
+def compute_sd(charges, configurations, n_processes=None):
+    """
+    Compute standard deviations for potential and electric field with optional parallelization.
+    
+    Args:
+        charges (numpy.ndarray): Charge values
+        configurations (list): List of configuration dictionaries
+        n_processes (int, optional): Number of processes for parallel computation
+        
+    Returns:
+        tuple: (vsd, esd) Standard deviations for potential and electric field
+    """
+    # Use parallel processing if requested and beneficial
+    if n_processes is not None and n_processes > 1 and len(configurations) > 1:
+        return _compute_sd_parallel(charges, configurations, n_processes)
+    else:
+        return _compute_sd_serial(charges, configurations)
+
+
+def _compute_sd_serial(charges, configurations):
+    """Serial version of compute_sd (original implementation)."""
+    eroh = 0.0
+    e_diff = 0.0
+
+    vroh = 0.0
+    v_diff = 0.0
+
+    for config in configurations:
+
+        ch_positions = config.get('qm_coordinates')
+        sr_positions = config.get('mm_coordinates')
+        sr_potential = config.get('electric_potential')
+        sr_electric_field = config.get('electric_field')
+
+        ediff = compute_diff_electric_field(charges, ch_positions, 
+                                        sr_positions, sr_electric_field)
+        vdiff = compute_diff_potential(charges, ch_positions, 
+                                    sr_positions, sr_potential)
+        
+        eroh += ediff[0]
+        e_diff += ediff[1]
+        vroh += vdiff[0]
+        v_diff += vdiff[1]
+
+    vsd = np.sqrt(v_diff/vroh)
+    esd = np.sqrt(e_diff/eroh)
+        
+    return vsd, esd
+
+
+def _process_single_config_sd(args):
+    """
+    Process a single configuration for parallel SD computation.
+    
+    Args:
+        args: Tuple containing (charges, config)
+        
+    Returns:
+        tuple: (eroh, e_diff, vroh, v_diff)
+    """
+    charges, config = args
+    
+    ch_positions = config.get('qm_coordinates')
+    sr_positions = config.get('mm_coordinates')
+    sr_potential = config.get('electric_potential')
+    sr_electric_field = config.get('electric_field')
+
+    ediff = compute_diff_electric_field(charges, ch_positions, 
+                                    sr_positions, sr_electric_field)
+    vdiff = compute_diff_potential(charges, ch_positions, 
+                                sr_positions, sr_potential)
+    
+    return ediff[0], ediff[1], vdiff[0], vdiff[1]
+
+
+def _compute_sd_parallel(charges, configurations, n_processes):
+    """Parallel version of compute_sd using multiprocessing."""
+    # Prepare arguments for parallel processing
+    args_list = [(charges, config) for config in configurations]
+    
+    # Use multiprocessing pool
+    with mp.Pool(processes=n_processes) as pool:
+        results = pool.map(_process_single_config_sd, args_list)
+    
+    # Sum up results from all processes
+    eroh = 0.0
+    e_diff = 0.0
+    vroh = 0.0
+    v_diff = 0.0
+    
+    for result in results:
+        eroh += result[0]
+        e_diff += result[1]
+        vroh += result[2]
+        v_diff += result[3]
+
+    vsd = np.sqrt(v_diff/vroh)
+    esd = np.sqrt(e_diff/eroh)
+        
+    return vsd, esd
+
+def compute_infulence_mat(configurations, wv, we, 
+                          eq_map, optimize_charges, n_processes=None):
+    """
+    Compute influence matrix with optional parallelization.
+    
+    Args:
+        configurations (list): List of configuration dictionaries
+        wv (float): Potential weight
+        we (float): Electric field weight
+        eq_map (dict): Equivalent atom mapping
+        optimize_charges (list): List of charges to optimize
+        n_processes (int, optional): Number of processes for parallel computation
+        
+    Returns:
+        numpy.ndarray: Influence matrix
+    """
+    # Use parallel processing if requested and beneficial
+    if n_processes is not None and n_processes > 1 and len(configurations) > 1:
+        return _compute_infulence_mat_parallel(configurations, wv, we, eq_map, optimize_charges, n_processes)
+    else:
+        return _compute_infulence_mat_serial(configurations, wv, we, eq_map, optimize_charges)
+
+
+def _compute_infulence_mat_serial(configurations, wv, we, eq_map, optimize_charges):
+    """Serial version of compute_infulence_mat (original implementation)."""
+    infulence_mat = []
+    for config in configurations:
+        single_inful = infulence_mat_single(config, wv, we, eq_map, optimize_charges)
+        infulence_mat.append(single_inful.T)
+    infulence_mat = np.vstack(infulence_mat)
+
+    return infulence_mat
+
+
+def _process_single_influence_mat(args):
+    """
+    Process a single configuration for parallel influence matrix computation.
+    
+    Args:
+        args: Tuple containing (config, wv, we, eq_map, optimize_charges)
+        
+    Returns:
+        numpy.ndarray: Single influence matrix
+    """
+    config, wv, we, eq_map, optimize_charges = args
+    single_inful = infulence_mat_single(config, wv, we, eq_map, optimize_charges)
+    return single_inful.T
+
+
+def _compute_infulence_mat_parallel(configurations, wv, we, eq_map, optimize_charges, n_processes):
+    """Parallel version of compute_infulence_mat using multiprocessing."""
+    # Prepare arguments for parallel processing
+    args_list = [(config, wv, we, eq_map, optimize_charges) for config in configurations]
+    
+    # Use multiprocessing pool
+    with mp.Pool(processes=n_processes) as pool:
+        results = pool.map(_process_single_influence_mat, args_list)
+    
+    # Stack results
+    infulence_mat = np.vstack(results)
+    
+    return infulence_mat
+
+def infulence_mat_single(config, wv, we, eq_map, optimize_charges):
+
+    columns = []
+    qm_coordinates = config['qm_coordinates']
+    mm_coordinates = config['mm_coordinates']
+    for i in range(qm_coordinates.shape[0]):
+        col = []
+        R_ij = mm_coordinates-qm_coordinates[i]
+        r_ij = np.linalg.norm(R_ij, axis=1)
+        dem = np.reshape(r_ij**3,(-1, 1))
+        poten_term = wv * 1/ r_ij
+        elec_term =  we * R_ij* 1/ dem
+        col.append(poten_term)
+        col.append(elec_term.flatten())
+        index = eq_map.get(i)
+        index = optimize_charges.index(index)
+        if (len(columns) - 1) < index:
+            columns.append(np.hstack(col))
+        else:
+            columns[index] += np.hstack(col)
+    return np.vstack(columns)
+
+def compute_target_mat(configurations, wv, we, wh, wq, q_total,
+                         q_restrain):
+
+    target_mat = []
+    for config in configurations:
+        potential = config['electric_potential']
+        elec_field = config['electric_field']
+        target_mat = np.hstack([target_mat, potential*wv, elec_field.flatten()*we])
+    target_mat = np.hstack([target_mat, q_restrain * wh, wq * q_total])
+    return target_mat
+
+
+def opt_dresp(configurations, wv, we, wh, wq, q_total,
+                        eq_map, n_processes=None, fixed_charge_indices=None,
+                        charge_group_constraints=None, weights_to_fix_charges=100000):
+    """
+    Optimize DRESP charges with optional parallelization and ability to fix certain charges.
+    
+    Args:
+        configurations (list): List of configuration dictionaries
+        wv (float): Potential weight
+        we (float): Electric field weight
+        wh (float): Restraint weight
+        wq (float): Total charge weight
+        q_total (float): Total charge constraint
+        eq_map (dict): Equivalent atom mapping
+        n_processes (int, optional): Number of processes for parallel computation
+        fixed_charge_indices (set, optional): Set of atom indices whose charges should be kept at original values
+        charge_group_constraints (list, optional): List of tuples (atom_indices, target_charge) where atom_indices 
+                                                 is a set of atom indices and target_charge is the desired sum of charges
+        
+    Returns:
+        numpy.ndarray: Optimized charges
+    """
+    
+    # Initialize fixed_charge_indices if None
+    if fixed_charge_indices is None:
+        fixed_charge_indices = set()
+    
+    # Initialize charge_group_constraints if None
+    if charge_group_constraints is None:
+        charge_group_constraints = []
+    
+    reference_charge = np.zeros(configurations[0]['qm_coordinates'].shape[0])
+    for config in configurations:
+        reference_charge += config['reference_charge']
+
+    reference_charge = reference_charge / len(configurations)
+
+    # Include all atoms in optimization, but use strong restraints for fixed charges
+    optimize_charges = []
+    for i in range(configurations[0]['qm_coordinates'].shape[0]):
+        idx = eq_map.get(i)
+        if idx not in optimize_charges:
+            optimize_charges.append(idx)
+
+    unique_charges = len(optimize_charges)
+    q_restrain = np.zeros(unique_charges)
+    nq_unique = np.zeros(unique_charges)
+    tot_charge_constraint= np.zeros(unique_charges)
+    
+    for i in range(reference_charge.shape[0]):
+        idx = eq_map.get(i)
+        idx = optimize_charges.index(idx)
+        q_restrain[idx] += reference_charge[i]
+        nq_unique[idx] += 1
+        tot_charge_constraint[idx] += wq
+    
+    # Fix division by zero: only divide where nq_unique > 0
+    mask = nq_unique > 0
+    q_restrain[mask] = q_restrain[mask] / nq_unique[mask]
+    q_restrain = q_restrain * nq_unique 
+
+    target_mat = compute_target_mat(configurations, wv, we, wh, wq, q_total,
+                        q_restrain)
+
+    
+    restrain_mat = np.diagflat(nq_unique) * wh
+
+    infulence_mat = compute_infulence_mat(configurations, wv, we, 
+                          eq_map, optimize_charges, n_processes)
+    
+    # Add charge group constraints
+    charge_group_constraint_rows = []
+    charge_group_targets = []
+    
+    for atom_indices, target_charge in charge_group_constraints:
+        constraint_row = np.zeros(unique_charges)
+        for atom_idx in atom_indices:
+            if atom_idx < len(reference_charge):
+                eq_idx = eq_map.get(atom_idx)
+                if eq_idx in optimize_charges:
+                    constraint_idx = optimize_charges.index(eq_idx)
+                    constraint_row[constraint_idx] += wq
+        
+        charge_group_constraint_rows.append(constraint_row)
+        charge_group_targets.append(target_charge * wq)
+    
+    # Add very strong restraints for fixed charges to force them to reference values
+    fixed_charge_restraints = np.zeros((len(fixed_charge_indices), unique_charges))
+    fixed_charge_targets = []
+    
+    for i, atom_idx in enumerate(fixed_charge_indices):
+        if atom_idx < len(reference_charge):
+            constraint_row = np.zeros(unique_charges)
+            eq_idx = eq_map.get(atom_idx)
+            if eq_idx in optimize_charges:
+                constraint_idx = optimize_charges.index(eq_idx)
+                # Use very strong restraint weight (1000 * wh) to effectively fix the charge
+                constraint_row[constraint_idx] = weights_to_fix_charges
+            fixed_charge_restraints[i] = constraint_row
+            fixed_charge_targets.append(reference_charge[atom_idx] * weights_to_fix_charges)
+    
+    # Stack all constraints
+    all_constraints = [restrain_mat, tot_charge_constraint]
+    if charge_group_constraint_rows:
+        all_constraints.extend(charge_group_constraint_rows)
+    if len(fixed_charge_restraints) > 0:
+        all_constraints.append(fixed_charge_restraints)
+    
+    infulence_mat = np.vstack([infulence_mat] + all_constraints)
+    
+    # Add targets to target matrix
+    additional_targets = charge_group_targets + fixed_charge_targets
+    if additional_targets:
+        target_mat = np.hstack([target_mat] + additional_targets)
+
+    res = np.linalg.lstsq(infulence_mat,target_mat, rcond=None)
+
+    full_optimize_charges = np.zeros(configurations[0]['qm_coordinates'].shape[0])
+    for i in range(configurations[0]['qm_coordinates'].shape[0]):
+        idx = eq_map.get(i)
+        idx = optimize_charges.index(idx)
+        full_optimize_charges[i] = res[0][idx]
+
+    return full_optimize_charges

mimicpy/force_matching/nonbonded_forces.py

@@ -0,0 +1,32 @@
+import numpy as np
+from ..utils.constants import kjmolnm_to_au, nm_to_au
+
+def get_qm_gmx_forces(mda, idx_frame, qm_atoms):
+    """
+    Get GROMACS forces and positions for QM atoms
+    
+    Args:
+        mda: MDAnalysis Universe object
+        idx_frame (int): Frame index
+        qm_atoms: Set or list of GROMACS atom indices (1-based) including QM atoms
+        
+    Returns:
+        tuple: (forces, positions) where forces and positions are numpy arrays
+    """
+    # mda.trajectory[0]
+    # if mda.trajectory[0].time == 0.0:
+    #     idx_frame += 1
+    mda.trajectory[idx_frame]
+    forces = [] 
+    positions = []
+    
+    # Convert qm_atoms to set for efficient lookup
+    if not isinstance(qm_atoms, set):
+        qm_atoms = set(qm_atoms)
+    
+    for atom in mda.atoms:
+        if atom.id+1 in qm_atoms:
+            forces.append(atom.force*10*kjmolnm_to_au)
+            positions.append(atom.position * 0.1 * nm_to_au)
+
+    return np.array(forces), np.array(positions)