print formating

TIES_MD/TIES.py

@@ -19,11 +19,7 @@
 __license__ = "LGPL"
 
 from .lambdas import Lambdas
-try:
-    from .alch import AlchSys, System_ID, simulate_system
-    openmm_av = True
-except ImportError:
-    openmm_av = False
+from .alch import AlchSys, System_ID, simulate_system
 
 try:
     from openmm import unit, Vec3
@@ -41,7 +37,7 @@
 from pathlib import Path
 
 import importlib.resources as pkg_resources
-#split means we will explitlly deal with reps in the submission
+#split means we will explicitly deal with reps in the submission
 #for non split namd or TIES_MD will handle the parallisim of replicas
 from .eng_scripts import namd_sub, namd_sub_split, openmm_sub_split, cfg_scripts
 
@@ -64,6 +60,13 @@ class TIES(object):
     def __init__(self, cwd, exp_name, run_type='class', devices=None, node_id=None, windows_mask=None, periodic=True,
                  lam=None, platform='CUDA', **kwargs):
         nice_print('TIES')
+        print('Wade, A.D., et al. 2022. Alchemical Free Energy Estimators and Molecular Dynamics Engines:'
+              ' Accuracy, Precision, and Reproducibility. Journal of chemical theory and computation, 18(6), pp.3972-3987.\n')
+
+        print('If you use OpenMM with this software please cite:')
+        print('Chodera, J., et al. 2022. choderalab/openmmtools: 0.21.4. Zenodo. doi: 10.5281/zenodo.6625266.')
+        print('Eastman, P., et al. 2017. OpenMM 7: Rapid development of high performance algorithms'
+              ' for molecular dynamics. PLoS computational biology, 13(7), p.e1005659.\n')
 
         #check all the config file args we need are present
         args_list = ['engine', 'temperature', 'pressure', 'sampling_per_window', 'equili_per_window', 'methods',
@@ -90,14 +93,15 @@ def __init__(self, cwd, exp_name, run_type='class', devices=None, node_id=None,
                          'box_type']
 
         #Iterate over our args_dict to set attributes of class to values in dict
-        print('Read arguments from file...')
+        print('Read arguments from file:')
         for k, v in kwargs.items():
             if k in api_sensitive:
                 full_k = '_'+k
             else:
                 full_k = k
             print('{} = {}'.format(k, v))
             setattr(self, full_k, v)
+        print('')
 
         #set any nonexistant optional args to None
         for option in optional_args:
@@ -168,12 +172,14 @@ def __init__(self, cwd, exp_name, run_type='class', devices=None, node_id=None,
                                                                 self.global_lambdas, self.total_reps, self.exp_name)
 
         # build schedule for lambdas do this last so passed lam can overwrite if desired
+        print('Lambda schedule:')
         if lam is None:
             # build lambda schedule
             self.lam = Lambdas(self.elec_edges, self._ster_edges, self.global_lambdas)
         else:
             self.lam = lam
             print('Using custom lambda schedule all schedule options will be ignored')
+        print('')
 
         #Perform a job
         if self.run_type == 'run':
@@ -566,6 +572,7 @@ def run(self):
         pool = Pool(processes=len(self.devices))
 
         tic = time.perf_counter()
+        nice_print('Beginning simulations')
         simulated = pool.map(func, system_ids) ### MAIN MD
         toc = time.perf_counter()
         total_simulation_time = (toc - tic)
@@ -585,7 +592,7 @@ def run(self):
         #The differnce between the time reported here and the time in the logs is from alchemy
         print('Simulation speed per GPU = {} ns/day'.format(speed))
 
-        nice_print('Finished simulation')
+        nice_print('Finished simulations')
 
     def write_analysis_cfg(self):
         '''

TIES_MD/alch.py

@@ -139,19 +139,22 @@ def __init__(self, cwd, name, temperature, pressure, constraint_file, constraint
             disappear_idxs = disappear_idxs+appear_idxs
             appear_idxs = []
 
+        print('Alchemical indexes:')
         print('Appearing atoms {}'.format(appear_idxs))
-        print('Disappearing atoms {}'.format(disappear_idxs))
+        print('Disappearing atoms {}\n'.format(disappear_idxs))
 
         self.appearing_idxs = appear_idxs
 
         print('Default box vectors were:')
-        print(system.getDefaultPeriodicBoxVectors())
+        for axis in system.getDefaultPeriodicBoxVectors():
+            print(axis)
 
         print('Switching to explicit box vectors:')
         system.setDefaultPeriodicBoxVectors(*basis_vectors)
         self.PBV = system.getDefaultPeriodicBoxVectors()
         for axis in self.PBV:
             print(axis)
+        print('')
 
         if debug:
             AlchSys.debug_force(self, system)
@@ -190,7 +193,7 @@ def __init__(self, cwd, name, temperature, pressure, constraint_file, constraint
                 for i in range(force.getNumGlobalParameters()):
                     param_name = force.getGlobalParameterName(i)
                     if 'lambda_sterics' in param_name:
-                        print('Adding derivative for {} to a CustomNonbondedForce or CustomBondForce'.format(param_name))
+                        #print('Adding derivative for {} to a CustomNonbondedForce or CustomBondForce'.format(param_name))
                         force.addEnergyParameterDerivative(param_name)
                         continue
 
@@ -209,6 +212,7 @@ def __init__(self, cwd, name, temperature, pressure, constraint_file, constraint
         self.NVT_compound_state = openmmtools.states.CompoundThermodynamicState(thermodynamic_state=ts,
                                                                             composable_states=composable_states)
         #Turn off intersect angle, bond, torsions
+        print('Cleaning system:')
         if len(intersect_angles) > 0:
             print('Found angles {} straddling alchemical regions these will be turned off'.format(intersect_angles))
             self.NVT_compound_state.lambda_angles_appear = 0.0
@@ -217,6 +221,7 @@ def __init__(self, cwd, name, temperature, pressure, constraint_file, constraint
             self.NVT_compound_state.lambda_bonds_appear = 0.0
         if len(intersect_torsions) > 0:
             print('Found torsion {} straddling alchemical regions these will be removed'.format(intersect_torsions))
+        print('')
 
         self.NPT_alchemical_system = copy.deepcopy(self.NVT_alchemical_system)
         self.NPT_compound_state = copy.deepcopy(self.NVT_compound_state)
@@ -230,7 +235,7 @@ def __init__(self, cwd, name, temperature, pressure, constraint_file, constraint
 
         # Add constraints to NVT system
         if self.constraints is not None:
-            print('Constraining NVT system')
+            print('Constraining NVT system...\n')
             AlchSys.add_consraints(self)
 
     def amend_original_positions(self, positions):
@@ -836,7 +841,7 @@ def simulate_system(ids, alch_sys, Lam, mask, cwd, niter, equili_steps, steps_pe
         equili_files = list(glob.iglob(equili_file))
 
         load = False
-        print('Load flag for equilibrated data is {}'.format(load))
+        #print('Load flag for equilibrated data is {}'.format(load))
         if len(equili_files) > 0 and load:
             # load the equilibriated state from disk
             print('Loading equilibriated state from disk located here: {}'.format(equili_files[0]))
@@ -868,7 +873,9 @@ def simulate_system(ids, alch_sys, Lam, mask, cwd, niter, equili_steps, steps_pe
 
         #Production
         for iteration in range(niter):
-            print('Propagating iteration {}/{} in state {}/{}'.format(iteration + 1, niter, i+mask[0], all_states))
+            print('Propagating iteration {}/{} in window {}/{} for replica {}'.format(iteration + 1,
+                                                                                        niter, i+mask[0]+1,
+                                                                                        all_states, ids.node_id))
 
             # propogate system in current state
             NPT['sim'].step(steps_per_iter)
@@ -894,17 +901,18 @@ def simulate_system(ids, alch_sys, Lam, mask, cwd, niter, equili_steps, steps_pe
         remove_simulation_reporters(NPT['sim'])
 
     #Save results to disk
+    print('Saving results')
     for i, j in enumerate(Lam.str_lams[mask[0]: mask[1]]):
         if 'TI' in alch_sys.methods:
             file = os.path.join(cwd, 'LAMBDA_{}'.format(j), 'rep{}'.format(ids.node_id), 'results',
                                 'TI.npy')
-            print('Saving {} result to disk'.format(file))
+            #print('Saving {} result to disk'.format(file))
             np.save(file, grads[i, :, :])
 
         if 'FEP' in alch_sys.methods:
             file = os.path.join(cwd, 'LAMBDA_{}'.format(j), 'rep{}'.format(ids.node_id), 'results',
                                 'FEP.npy')
-            print('Saving {} result to disk'.format(file))
+            #print('Saving {} result to disk'.format(file))
             np.save(file, u_kln[i, :, :])
 
     #clean up