Multi dihedral allowed

bin/ff_assess

@@ -4,6 +4,7 @@ Main exe for scoring simulations against mapped trajectories
 """
 import argparse
 from glob import glob
+from pathlib import Path
 import numpy as np
 import MDAnalysis as mda
 import vermouth.forcefield
@@ -32,7 +33,7 @@ def report_interactions(h_score_dict, score_dict, exclude_outliers=False):
     buff = ''
 
     buff += " [ Interaction Distribution Report ]\n"
-    buff += f"   Overall Score : {{mean:.2f}} ± {{std:.2f}}\n\n"
+    buff += f"   Overall Score : {{mean:.3}} ± {{std:.3}}\n\n"
 
     buff += " Interaction Scores:\n"
     buff += " 0 - identical, 1 - no overlap\n\n"
@@ -49,7 +50,6 @@ def report_interactions(h_score_dict, score_dict, exclude_outliers=False):
     else:
         buff += '\n'
 
-    h_scores = []
     overall_scores = []
     for interaction_type in h_score_dict.keys():
         buff += f" {interaction_type} \n"
@@ -59,7 +59,6 @@ def report_interactions(h_score_dict, score_dict, exclude_outliers=False):
                 buff += f"\t{atom_group:20s}: {h_score:.2f} ({score:.2f})(excluded)\n"
             else:
                 buff += f"\t{atom_group:20s}: {h_score:.2f} ({score:.2f})\n"
-                h_scores.append(h_score)
                 overall_scores.append(score)
         buff += '\n'
 
@@ -89,7 +88,7 @@ def report_distances(score_matrix, molname, atom_names):
     score_matrix_report = np.column_stack([[''] + atom_names, score_matrix_report])
     buff = ''
     buff += f" [ Distance Distribution Report for {molname} ]\n "
-    buff += f" Overall Score : {{mean:.2f}} ± {{std:.2f}}\n\n"
+    buff += f" Overall Score : {{mean:.3}} ± {{std:.3}}\n\n"
     buff += f"   Max Score : {{max:.2f}}\n\n"
     buff += " Score guide:\n"
     buff += "   0.0-0.3 : good\n"
@@ -105,6 +104,7 @@ def report_distances(score_matrix, molname, atom_names):
             buff += f"{row[0]:>5} " + " ".join(f"{float(score):>5.2f}" for score in row[1:]) + "\n" # First element is the atom name, rest are scores
 
     off_diag = score_matrix[~np.eye(score_matrix.shape[0], dtype=bool)] # Exclude diagonal (0) for mean and std calculation
+    sem = np.std(off_diag, ddof=1) / np.sqrt(len(off_diag)) # Standard error of the mean
     printable = buff.format(molname=molname,
                             mean=np.mean(off_diag),
                             std=np.std(off_diag)/np.sqrt(off_diag.size),
@@ -120,6 +120,7 @@ def __main__():
     parser.add_argument('-i', type=str, dest="itp_files", help="itp file", nargs='*')
     parser.add_argument('-d', type=str, dest='reference',
                         help="Path to directory with reference distributions")
+    parser.add_argument('-pref', type=str, dest="prefix", help="common prefix to filename", default="")
     parser.add_argument('-plots', default=None, const=".", nargs="?",
                         dest="plots",
                         help="Make plots comparing distributions. Optionally provide a path (default: current dir)")
@@ -131,21 +132,31 @@ def __main__():
                         help='weight of the Hellinger distance in the distance score')
     parser.add_argument('-include-constraints', default=False, action='store_true', dest='include_constraints',
                         help='fully include constrained distances in the distance score score calculation')
-    parser.add_argument('-dists', default=False, action='store_true',
+    parser.add_argument('-dists', default=None, const=".", nargs="?",
                         dest="distribution_data",
-                        help="Save text files with time series and distribution data for interactions")
+                        help="Save text files with time series and distribution data for interactions. Optionally provide a path (default: current dir).")
 
     args = parser.parse_args()
 
     distribution_files = glob(f'{args.reference}/*distr.dat')
+    # create file map
+    file_map = {}
+    for f in distribution_files:
+        p = Path(f)
+        if p.name in file_map:
+            print(f"Warning: duplicate reference file name '{p.name}', using the first one.")
+            continue
+        file_map[p.name] = p
 
     # load trajectory
     if args.tprfile:
         u = mda.Universe(args.tprfile, args.trajfile, in_memory=True)
     else:
         u = mda.Universe(args.trajfile, in_memory=True)
 
-    plot_data = defaultdict(dict)
+    if args.plots:
+        plot_data = defaultdict(dict) # initialize only if needed
+
     # if itp file is provided use it
     if args.itp_files:
         ff = vermouth.forcefield.ForceField("dummy")
@@ -168,42 +179,63 @@ def __main__():
             interaction_groups, _, _ = interaction_mapper.get_interactions_group(molname)
             for inter_type in ['bonds', 'angles', 'dihedrals']:
                 for group_name, pair_idxs in interaction_groups[inter_type].items():
-                    distr, _ = interaction_distribution(u, inter_type, pair_idxs, save=args.distribution_data)
+                    distr, _ = interaction_distribution(u, inter_type, pair_idxs, group_name, args.prefix, save=args.distribution_data)
                     # calculate simulation distribution
                     probs = distr.T[1]
                     # read in reference distribution
+                    reference_name = f"{args.prefix}{group_name}_{inter_type}_distr.dat"
                     try:
-                        reference_data = np.loadtxt([i for i in distribution_files if group_name in i and inter_type in i][0])
-                    except IndexError:
-                        print(f"{group_name} file not found!")
+                        ref_file = file_map[reference_name]
+                    except KeyError:
+                        print(f"{reference_name} file not found!")
+                        if args.prefix == "":
+                            print("your prefix is set to the default value: consider changing it to the actual prefix.")
                         continue
+
+                    reference_data = np.genfromtxt(ref_file)
                     # calculate hellinger distance between simulated and reference distributions
                     h_score = hellinger(probs/probs.sum(), reference_data.T[1] / reference_data.T[1].sum())
                     h_score_dict[inter_type][group_name] = h_score
 
-                    score = calc_score(reference_data.T[1]/ reference_data.T[1].sum(), probs/probs.sum(), [args.hellinger_weight, 1-args.hellinger_weight], inter_type)
+                    score = calc_score(
+                        reference_data.T[1]/ reference_data.T[1].sum(), probs/probs.sum(), 
+                        [args.hellinger_weight, 1-args.hellinger_weight], inter_type
+                        )
                     score_dict[inter_type][group_name] = score
-
-                    plot_data[inter_type][group_name] = {"x": reference_data.T[0],
+                    if args.plot_data:
+                        plot_data[inter_type][group_name] = {"x": reference_data.T[0],
                                                          "Reference": reference_data.T[1],
                                                          'Simulated': probs}
         if args.plots:
             make_distribution_plot(plot_data, args.plot_data, name=f'{args.plots}/interaction_distributions_{molname}')
+
         report_interactions(h_score_dict, score_dict, args.exclude_outliers)
 
         # evaluation of distances
+        file_map = {}
+        for f in distribution_files:
+            p = Path(f)
+            if p.name in file_map:
+                print(f"Warning: duplicate reference file name '{p.name}', using the first one.")
+                continue
+            file_map[p.name] = p
+
         for molname, block in ff.blocks.items():
-            score_m, plot_data = score_matrix(molname, 
-                                   block, 
-                                   u, 
-                                   distribution_files,
-                                   args.hellinger_weight,
-                                   args.include_constraints)
-
-            atom_names = [name for _, name in block.nodes(data='atomname')]
+            score_m, plot_data = score_matrix(molname=molname, 
+                                   block=block, 
+                                   universe=u, 
+                                   file_map=file_map,
+                                   file_prefix=args.prefix,
+                                   hellinger_weight=args.hellinger_weight,
+                                   include_constrains=args.include_constraints)
+
+            atom_names = [f'{resid}_{name}' for (_, name), (_, resid) in zip(block.nodes(data='atomname'), block.nodes(data='resid'))]
+
             if args.plots:
                 make_matrix_plot(score_m, atom_names, name=f'{args.plots}/score_matrix_{molname}')
-                make_distances_distribution_plot(plot_data, atom_names, name=f'{args.plots}/distance_distributions_{molname}')
-            report_distances(score_m, molname, atom_names)
+                make_distances_distribution_plot(
+                    plot_data, atom_names, name=f'{args.plots}/distance_distributions_{molname}',save_plot_data=True
+                )
 
+            report_distances(score_m, molname, atom_names)
 __main__()

bin/ff_inter

@@ -19,6 +19,7 @@ Main exe for computing interactions.
 """
 import argparse
 import MDAnalysis as mda
+from tqdm import tqdm
 import vermouth.forcefield
 import sys
 import fast_forward
@@ -89,6 +90,8 @@ def __main__():
                         dest="dist_matrix", help="Save text files with time series and distribution data for all pairwise bead distances")
     parser.add_argument('-lincs', default=False, action="store_true",
                         dest="estimate_lincs", help="Estimate the LINCS order required for solving the constraints in the molecule(s)")
+    parser.add_argument('-multi-dihedral', default=False, action="store_true",
+                        dest="multi_dihedral", help="Fit multiple dihedral terms on top of each other and ignore penalty")
 
     args = parser.parse_args()
 
@@ -110,7 +113,8 @@ def __main__():
                                    max_dihedrals=args.max_dihedral,
                                    temperature=args.temperature,
                                    precision=args.precision,
-                                   dihedral_scaling=args.dihedral_scaling)
+                                   dihedral_scaling=args.dihedral_scaling,
+                                   disable_aic_penalty=args.multi_dihedral)
 
         interaction_mapper = ITPInteractionMapper(u, ff.blocks.values(), ff.blocks.keys())
         # loop over molecules
@@ -158,18 +162,23 @@ def __main__():
                        ff.blocks[molname],
                        command_used=' '.join(sys.argv),
                        )
-        if args.dist_matrix: # calculate and save the timeseries and distributions of all interactions
-            for molname, block in ff.blocks.items():
-                inter_type='distances'
-                interaction_groups = interaction_mapper.get_pairwise_interaction(molname)
-                for group_name, pair_idxs in interaction_groups[inter_type].items():
-                    if not args.distribution_data: # if no path provided, save to current dir, distribution data is always saved
-                        save='.'
-                    else:
-                        save=args.distribution_data
-                    distr = interaction_distribution(u, inter_type, pair_idxs, group_name,
-                                                    args.prefix, save=save)
-
+        if args.dist_matrix:
+            inter_type = 'distances'
+            total = sum(
+                len(interaction_mapper.get_pairwise_interaction(molname)[inter_type])
+                for molname in ff.blocks
+            )
+
+            with tqdm(total=total, desc='Calculating pairwise distances') as pbar:
+                for molname, block in ff.blocks.items():
+                    interaction_groups = interaction_mapper.get_pairwise_interaction(molname)
+                    for group_name, pair_idxs in interaction_groups[inter_type].items():
+                        save = args.distribution_data if args.distribution_data else '.'
+                        distr = interaction_distribution(
+                            u, inter_type, pair_idxs, group_name, args.prefix, save=save
+                        )
+                        pbar.set_postfix_str(molname)
+                        pbar.update(1)
         # constraint coupling analysis
         for molname, block in ff.blocks.items():
             report_constraint_coupling(u, block, args.estimate_lincs)

fast_forward/interaction_fit.py

@@ -96,7 +96,7 @@ class InteractionFitter:
     Class to fit interactions
     """
     def __init__(self, precision, temperature, constraint_converter,
-                 max_dihedrals, dihedral_scaling):
+                 max_dihedrals, dihedral_scaling, disable_aic_penalty=False):
         '''
 
         Parameters
@@ -109,6 +109,8 @@ def __init__(self, precision, temperature, constraint_converter,
             threshold above which to convert bonds to constraints
         max_dihedrals: int
             maximum number of dihedrals to fit proper dihedrals with
+        disable_aic_penalty: bool
+            if True, disable AIC penalty for dihedral fitting and use all available terms
         '''
         self.__dihedrals = None
         self.precision = precision
@@ -117,6 +119,7 @@ def __init__(self, precision, temperature, constraint_converter,
         self.constraint_converter = constraint_converter
         self.max_dihedrals = max_dihedrals
         self.dihedral_scaling = dihedral_scaling
+        self.disable_aic_penalty = disable_aic_penalty
         # this will store the interactions
         self.interactions_dict = defaultdict(list)
         self.fit_parameters = defaultdict(dict)
@@ -287,7 +290,11 @@ def _dihedrals_fitter(self, data, group_name):
 
         num_terms = len(best_params) // 3  # Each term has k, n, and x0
 
-        condition0 = best_aic < gaussian_result.aic
+        # If AIC penalty is disabled, always use proper dihedrals with all fitted terms
+        if self.disable_aic_penalty:
+            condition0 = True
+        else:
+            condition0 = best_aic < gaussian_result.aic
         condition1 = np.isclose(gaussian_result.params['sigma'].value, gaussian_result.params['sigma'].max)
 
         # compare the aic values to determine which type of dihedral we have

fast_forward/interaction_plots.py

@@ -29,7 +29,26 @@ def _plotter(data, atom_list, inter_type, ax):
     ax.set_xlabel(X_LABELS[inter_type])
 
 
-def _plotter_distance_distribution(data, ax):
+def _plotter_distance_distribution(data, ax, y_lower_threshold: float = 0.01):
+    """
+    Plot distance distribution curves and automatically zoom in on
+    regions where the signal exceeds a given fraction of its maximum.
+
+    Parameters
+    ----------
+    data : dict
+        Dictionary containing the distance distribution data.
+        Must include a key ``"x"`` for the shared x-axis values.
+        All other keys correspond to y-data series to be plotted and
+        must be array-like and of the same length as ``data["x"]``.
+    ax : matplotlib.axes.Axes
+        Axes object on which the curves will be plotted.
+    y_lower_threshold : float, optional
+        Fraction of each curve's maximum used to determine the
+        region of interest. Only x-values where ``y > y.max() * threshold``
+        are considered when computing the zoomed x-axis limits.
+        Default is ``0.01``
+    """
     cols = ['#6970E0', '#E06B69']
     needed_keys = [key for key in list(data.keys()) if key != 'x']
     x_min = 100
@@ -41,8 +60,12 @@ def _plotter_distance_distribution(data, ax):
                 data[key],
                 c = cols[idx],
                 label=key)
-        x_min = np.min([x_min, data['x'][np.min(np.nonzero(data[key]))]]) # find the minimum x value in the data
-        x_max = np.max([x_max, data['x'][np.max(np.nonzero(data[key]))]]) # find the maximum x value in the data
+        threshold = np.max(data[key]) * y_lower_threshold
+        significant_indices = np.where(data[key] > threshold)[0]
+
+        if significant_indices.size > 0: # Only update if we found significant data
+            x_min = np.min([x_min, data['x'][np.min(significant_indices)]])
+            x_max = np.max([x_max, data['x'][np.max(significant_indices)]])
     ax.yaxis.set_ticks([])
     ax.set_xlim(x_min - x_pad, x_max + x_pad)
 
@@ -154,13 +177,16 @@ def make_distances_distribution_plot(plot_data, atom_names, save_plot_data=False
     if not axarr:
         fig ,axarr = plt.subplots(natoms-1,natoms-1,figsize=(natoms*2,natoms),gridspec_kw={'wspace':0.05,'hspace':0.4})
 
+    parsed_names = [name.split("_") for name in atom_names]
     for i in range(natoms-1):
+        resid1, name1 = parsed_names[i]
         for j in range(1,natoms):
+            resid2, name2 = parsed_names[j]
             ax = axarr[i, j-1]
             if i < j: # plot only upper triangle of the matrix
-                atoms = f'{atom_names[i]}_{atom_names[j]}'
-                if atoms in plot_data['distances']:
-                    _plotter_distance_distribution(plot_data['distances'][atoms], ax)
+                atoms_key = f'{resid1}_{resid2}_{name1}_{name2}'
+                if atoms_key in plot_data['distances']:
+                    _plotter_distance_distribution(plot_data['distances'][atoms_key], ax)
             else:
                 fig.delaxes(ax) # remove lower triangle of the matrix
 

fast_forward/itp_to_ag.py

@@ -97,7 +97,7 @@ def get_pairwise_interaction(self, molname):
         for node1, name1 in block.nodes(data='atomname'):
             for node2, name2 in list(block.nodes(data='atomname'))[node1+1:]:
                 atoms = np.array([node1, node2])
-                group = f'{name1}_{name2}' # naming convention with node1 < node2
+                group = f'{block.nodes[node1]["resid"]}_{name1}_{block.nodes[node2]["resid"]}_{name2}' # naming convention with node1 < node2
                 indices = find_mol_indices(self.universe, atoms, molname)
                 old_indices = indices_dict['distances'].get(group, [])
                 indices_dict['distances'][group] = indices + old_indices