From 3601e1a655c2c0cc787bb2b48fbe7c73cedc0145 Mon Sep 17 00:00:00 2001 From: sampurna Date: Mon, 25 Mar 2024 21:47:35 +0530 Subject: [PATCH] Added tqdm bar for: pca.transform() --- package/MDAnalysis/analysis/pca.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/package/MDAnalysis/analysis/pca.py b/package/MDAnalysis/analysis/pca.py index 4ed5486e4c..c1cd216875 100644 --- a/package/MDAnalysis/analysis/pca.py +++ b/package/MDAnalysis/analysis/pca.py @@ -121,11 +121,12 @@ from MDAnalysis import Universe from MDAnalysis.analysis.align import _fit_to -from MDAnalysis.lib.log import ProgressBar +from MDAnalysis.lib.log import ProgressBar, ProgressBar_Dual from ..lib import util from ..due import due, Doi from .base import AnalysisBase +from tqdm.auto import tqdm class PCA(AnalysisBase): @@ -355,12 +356,12 @@ def n_components(self, n): n = len(self._variance) self.results.variance = self._variance[:n] self.results.cumulated_variance = (np.cumsum(self._variance) / - np.sum(self._variance))[:n] + np.sum(self._variance))[:n] self.results.p_components = self._p_components[:, :n] self._n_components = n def transform(self, atomgroup, n_components=None, start=None, stop=None, - step=None): + step=None, verbose=None): """Apply the dimensionality reduction on a trajectory Parameters @@ -382,6 +383,7 @@ def transform(self, atomgroup, n_components=None, start=None, stop=None, Include every `step` frames in the PCA transform. If set to ``None`` (the default) then every frame is analyzed (i.e., same as ``step=1``). + verbose: Boolean,optional Returns ------- @@ -398,7 +400,7 @@ def transform(self, atomgroup, n_components=None, start=None, stop=None, if isinstance(atomgroup, Universe): atomgroup = atomgroup.atoms - if(self._n_atoms != atomgroup.n_atoms): + if (self._n_atoms != atomgroup.n_atoms): raise ValueError('PCA has been fit for' '{} atoms. Your atomgroup' 'has {} atoms'.format(self._n_atoms, @@ -415,7 +417,15 @@ def transform(self, atomgroup, n_components=None, start=None, stop=None, dot = np.zeros((n_frames, dim)) - for i, ts in enumerate(traj[start:stop:step]): + if verbose == True: + self.disabled = False + else: + self.disabled = True + + # for i, ts in tqdm(enumerate(traj[start:stop:step]), desc="Transform progress"): + for i, ts in tqdm(enumerate(traj[start:stop:step]), disable=self.disabled, desc="Transform progress"): + # print("Iteration", i, "Timestep",ts) + # verbose=self._verbose, desc="Transform progress"): xyz = atomgroup.positions.ravel() - self._xmean dot[i] = np.dot(xyz, self._p_components[:, :dim]) @@ -561,7 +571,7 @@ def project_single_frame(self, components=None, group=None, anchor=None): for res in group.residues: # n_common is the number of pca atoms in a residue n_common = pca_res_counts[np.where( - pca_res_indices == res.resindex)][0] + pca_res_indices == res.resindex)][0] non_pca_atoms = np.append(non_pca_atoms, res.atoms.n_atoms - n_common) # index_extrapolate records the anchor number for each non-PCA atom