add threadpoolctl to clustenm app

jamesmkrieger · Oct 20, 2023 · 6b644aa · 6b644aa
1 parent 5acb01a
commit 6b644aa
Showing 1 changed file with 35 additions and 13 deletions.
diff --git a/prody/apps/prody_apps/prody_clustenm.py b/prody/apps/prody_apps/prody_clustenm.py
@@ -141,19 +141,41 @@ def prody_clustenm(pdb, **kwargs):
         except TypeError:
             raise TypeError("Please provide cutoff as a float or equation using math")
 
-    ens = prody.ClustENM(pdb.getTitle())
-    ens.setAtoms(select)
-    ens.run(n_gens=ngens, n_modes=nmodes,
-            n_confs=nconfs, rmsd=eval(rmsd),
-            cutoff=cutoff, gamma=gamma,
-            maxclust=eval(maxclust), threshold=eval(threshold),
-            solvent=solvent, force_field=eval(forcefield),
-            sim=sim, temp=temp, t_steps_i=t_steps_i,
-            t_steps_g=eval(t_steps_g),
-            outlier=outlier, mzscore=mzscore,
-            sparse=sparse, kdtree=kdtree, turbo=turbo,
-            parallel=parallel, fitmap=fitmap, 
-            fit_resolution=fit_resolution, **kwargs)
+    nproc = kwargs.get('nproc')
+    if nproc:
+        try:
+            from threadpoolctl import threadpool_limits
+        except ImportError:
+            raise ImportError('Please install threadpoolctl to control threads')
+
+        with threadpool_limits(limits=nproc, user_api="blas"):
+            ens = prody.ClustENM(pdb.getTitle())
+            ens.setAtoms(select)
+            ens.run(n_gens=ngens, n_modes=nmodes,
+                    n_confs=nconfs, rmsd=eval(rmsd),
+                    cutoff=cutoff, gamma=gamma,
+                    maxclust=eval(maxclust), threshold=eval(threshold),
+                    solvent=solvent, force_field=eval(forcefield),
+                    sim=sim, temp=temp, t_steps_i=t_steps_i,
+                    t_steps_g=eval(t_steps_g),
+                    outlier=outlier, mzscore=mzscore,
+                    sparse=sparse, kdtree=kdtree, turbo=turbo,
+                    parallel=parallel, fitmap=fitmap,
+                    fit_resolution=fit_resolution, **kwargs)
+    else:
+        ens = prody.ClustENM(pdb.getTitle())
+        ens.setAtoms(select)
+        ens.run(n_gens=ngens, n_modes=nmodes,
+                n_confs=nconfs, rmsd=eval(rmsd),
+                cutoff=cutoff, gamma=gamma,
+                maxclust=eval(maxclust), threshold=eval(threshold),
+                solvent=solvent, force_field=eval(forcefield),
+                sim=sim, temp=temp, t_steps_i=t_steps_i,
+                t_steps_g=eval(t_steps_g),
+                outlier=outlier, mzscore=mzscore,
+                sparse=sparse, kdtree=kdtree, turbo=turbo,
+                parallel=parallel, fitmap=fitmap,
+                fit_resolution=fit_resolution, **kwargs)
 
     single = not kwargs.pop('multiple')
     outname = join(outdir, prefix)