Repository: openmm/openmm — pull request #4671 (cpp)
Title: Increase describeNextReport readability and allow requesting more values

Fixes #4663 
Also harmonizes _generate_reports with #4437 

describeNextReport can now return a dict instead of a tuple with the format
`return {'steps':steps, 'periodic':False, 'include':['positions', 'forces']}`

The old tuple format is still supported.

This makes it mainly possible to also request parameterDerivatives for TI calculations without having to get the state for that reporter separately.
It also makes the output of describeNextReport much more readable than the previous tuple.

---
# Example
Following is an example reporter utilizing the new format

```python
class dUdLReporter(object):
    """
    Reporter for printing out current dUdL values.

    Calculate derivative needs to be set for this to work.

    Based on openMM source code for StateDataReporter, and openMM user guide http://docs.openmm.org/latest/userguide/application/04_advanced_sim_examples.html#extracting-and-reporting-forces-and-other-data

    Args:
        object (_type_): _description_
    """
    def __init__(self, file, reportInterval: int, derivativeNames: list, append: bool =False, timeUnit: unit = picosecond):
        """
        Create reporter for writing out energies of individual force groups

        Args:
            file (str or stdout): where to print the values
            reportInterval (int): report every n steps
            derivativeNames (list, optional): List of str defining the names of derivatives to report.
            append (bool, optional): append to file. Defaults to False.
            timeUnit (mm.unit, optional): time unit to use. Defaults to unit.picosecond.
        """

        self._append = append
        self._derivativeNames = derivativeNames
        self._reportInterval = reportInterval
        self._timeUnit = timeUnit

        # Open the file if it is a file, otherwise write directly to it (probably stdout)
        self._openedFile = isinstance(file, str)
        if self._openedFile:
            self._out = open(file, 'a' if append else 'w')
        else:
            self._out = file

        self._hasInitialized = False

    def __del__(self):
        # Close the file upon deletion
        if self._openedFile:
            self._out.close()

    def describeNextReport(self, simulation):
        # returns in how many steps the next report is due, and which kind of data it needs
        steps = self._reportInterval - simulation.currentStep%self._reportInterval
        return {'steps': steps, 'periodic': None, 'include':['parameterDerivatives']}

    def report(self, simulation, state=None):
        # Generates the report. Simulation and state is passed through by the openMM engine.

        # print header if it is the first time this object has been called and is not set to append.
        if not self._hasInitialized:
            self._hasInitialized = True
            if not self._append:
                header = ','.join(self._derivativeNames)
                header = "time," + header

                print(header, file=self._out)
                try:
                    self._out.flush()
                except AttributeError:
                    pass

        derivatives = []
        if state is None:
            state = simulation.context.getState(getParameterDerivatives=True)
        
        for derivativeName in self._derivativeNames:
            derivatives.append(f"{state.getEnergyParameterDerivatives()[derivativeName]: .8f}")

        values = ','.join(derivatives)
        values = f"{state.getTime().value_in_unit(self._timeUnit): .2f},{values}"

        print(values, file=self._out)
        try:
            self._out.flush()
        except AttributeError:
            pass
```

---
# Todo

- [x] write documentation

---
# Additional thoughts

It might also make sense to include a reporter similar as the one above.

Should existing reporters be adjusted to match the proposed format?


Resolved issues:
title: Allow to request all parameters context.getState hast to offer from a reporters describeNextReport() function
body: At the moment, custom reporters are not able to request parameter derivatives or integrator parameters for their next report.

https://github.com/openmm/openmm/blob/fed44c579e9ac3136fe38056968a9c6dc1b6931e/wrappers/python/openmm/app/simulation.py#L254-L272

Currently, to get parameter derivatives using a custom reporter, the state needs to be fetched separately.

I can make the changes and create a pull request if desired.

## IMPORTANT INSTRUCTIONS

1. If you edit a python file, you MUST run 'make PythonInstall' inside the build directory
2. If you edit a c++ file, you MUST run 'make -j4 && make install" inside the build directory

Hints:

Fix patch:
diff --git a/docs-source/usersguide/application/04_advanced_sim_examples.rst b/docs-source/usersguide/application/04_advanced_sim_examples.rst
index 88e550caf3..88f49959e3 100644
--- a/docs-source/usersguide/application/04_advanced_sim_examples.rst
+++ b/docs-source/usersguide/application/04_advanced_sim_examples.rst
@@ -137,7 +137,7 @@ Here is the definition of the :class:`ForceReporter` class:
 
             def describeNextReport(self, simulation):
                 steps = self._reportInterval - simulation.currentStep%self._reportInterval
-                return (steps, False, False, True, False, None)
+                return = {'steps': steps, 'periodic': None, 'include':['forces']}
 
             def report(self, simulation, state):
                 forces = state.getForces().value_in_unit(kilojoules/mole/nanometer)
@@ -157,18 +157,14 @@ We then have two methods that every reporter must implement:
 :meth:`describeNextReport()` and :meth:`report()`.  A Simulation object
 periodically calls :meth:`describeNextReport()` on each of its reporters to
 find out when that reporter will next generate a report, and what information
-will be needed to generate it.  The return value should be a six element tuple,
-whose elements are as follows:
+will be needed to generate it.  The return value should be a dictionary with the following content:
 
-* The number of time steps until the next report.  We calculate this as
+* :code:`steps` (int): The number of time steps until the next report.  We calculate this as
   *(report interval)*\ -\ *(current step)*\ %\ *(report interval)*\ .  For
   example, if we want a report every 100 steps and the simulation is currently on
   step 530, we will return 100-(530%100) = 70.
-* Whether the next report will need particle positions.
-* Whether the next report will need particle velocities.
-* Whether the next report will need forces.
-* Whether the next report will need energies.
-* Whether the positions should be wrapped to the periodic box.  If None, it will
+* :code:`include` (list of strings): The types of information that need to be included in the next report. The values in the list correspond to arguments to :meth:`Context.getState()`. Allowed values include :code:`'positions'`, :code:`'velocities'`, :code:`'forces'`, :code:`'energy'`, :code:`'parameters'`, :code:`'parameterDerivatives'`, and :code:`'integratorParameters'`.
+* :code:`periodic` (bool, optional): Whether the positions should be wrapped to the periodic box.  If None or not set, it will
   automatically decide whether to wrap positions based on whether the System uses
   periodic boundary conditions.
 
diff --git a/wrappers/python/openmm/app/checkpointreporter.py b/wrappers/python/openmm/app/checkpointreporter.py
index 9177d2c17e..ebff7e8645 100644
--- a/wrappers/python/openmm/app/checkpointreporter.py
+++ b/wrappers/python/openmm/app/checkpointreporter.py
@@ -110,14 +110,11 @@ def describeNextReport(self, simulation):
 
         Returns
         -------
-        tuple
-            A five element tuple. The first element is the number of steps
-            until the next report. The remaining elements specify whether
-            that report will require positions, velocities, forces, and
-            energies respectively.
+        dict
+            A dictionary describing the required information for the next report
         """
         steps = self._reportInterval - simulation.currentStep%self._reportInterval
-        return (steps, False, False, False, False)
+        return {'steps':steps, 'periodic':None, 'include':[]}
 
     def report(self, simulation, state):
         """Generate a report.
diff --git a/wrappers/python/openmm/app/dcdreporter.py b/wrappers/python/openmm/app/dcdreporter.py
index f9020a6a82..09b1d9439a 100644
--- a/wrappers/python/openmm/app/dcdreporter.py
+++ b/wrappers/python/openmm/app/dcdreporter.py
@@ -78,15 +78,11 @@ def describeNextReport(self, simulation):
 
         Returns
         -------
-        tuple
-            A six element tuple. The first element is the number of steps
-            until the next report. The next four elements specify whether
-            that report will require positions, velocities, forces, and
-            energies respectively.  The final element specifies whether
-            positions should be wrapped to lie in a single periodic box.
+        dict
+            A dictionary describing the required information for the next report
         """
         steps = self._reportInterval - simulation.currentStep%self._reportInterval
-        return (steps, True, False, False, False, self._enforcePeriodicBox)
+        return {'steps':steps, 'periodic':self._enforcePeriodicBox, 'include':['positions']}
 
     def report(self, simulation, state):
         """Generate a report.
diff --git a/wrappers/python/openmm/app/internal/debuggingreporter.py b/wrappers/python/openmm/app/internal/debuggingreporter.py
index 2889b4dfd8..05eebf3a4c 100644
--- a/wrappers/python/openmm/app/internal/debuggingreporter.py
+++ b/wrappers/python/openmm/app/internal/debuggingreporter.py
@@ -67,14 +67,10 @@ def describeNextReport(self, simulation):
 
         Returns
         -------
-        tuple
-            A six element tuple. The first element is the number of steps
-            until the next report. The next four elements specify whether
-            that report will require positions, velocities, forces, and
-            energies respectively.  The final element specifies whether
-            positions should be wrapped to lie in a single periodic box.
+        dict
+            A dictionary describing the required information for the next report
         """
-        return (1, True, True, True, False, False)
+        return {'steps':1, 'periodic':False, 'include':['positions', 'velocities', 'forces']}
 
     def report(self, simulation, state):
         """Generate a report.
diff --git a/wrappers/python/openmm/app/pdbreporter.py b/wrappers/python/openmm/app/pdbreporter.py
index d9b6c7593c..256ed36f17 100644
--- a/wrappers/python/openmm/app/pdbreporter.py
+++ b/wrappers/python/openmm/app/pdbreporter.py
@@ -77,15 +77,11 @@ def describeNextReport(self, simulation):
 
         Returns
         -------
-        tuple
-            A six element tuple. The first element is the number of steps
-            until the next report. The next four elements specify whether
-            that report will require positions, velocities, forces, and
-            energies respectively.  The final element specifies whether
-            positions should be wrapped to lie in a single periodic box.
+        dict
+            A dictionary describing the required information for the next report
         """
         steps = self._reportInterval - simulation.currentStep%self._reportInterval
-        return (steps, True, False, False, False, self._enforcePeriodicBox)
+        return {'steps':steps, 'periodic':self._enforcePeriodicBox, 'include':['positions']}
 
     def report(self, simulation, state):
         """Generate a report.
diff --git a/wrappers/python/openmm/app/simulatedtempering.py b/wrappers/python/openmm/app/simulatedtempering.py
index 1451b69b91..b4a23fafa4 100644
--- a/wrappers/python/openmm/app/simulatedtempering.py
+++ b/wrappers/python/openmm/app/simulatedtempering.py
@@ -176,7 +176,10 @@ def describeNextReport(self, simulation):
                 steps2 = st.reportInterval - simulation.currentStep%st.reportInterval
                 steps = min(steps1, steps2)
                 isUpdateAttempt = (steps1 == steps)
-                return (steps, False, isUpdateAttempt, False, isUpdateAttempt)
+                if isUpdateAttempt:
+                    return {'steps': steps, 'periodic':None, 'include':['velocities', 'energy']}
+                else:
+                    return {'steps': steps, 'periodic':None, 'include':[]}
 
             def report(self, simulation, state):
                 st = self.st
diff --git a/wrappers/python/openmm/app/simulation.py b/wrappers/python/openmm/app/simulation.py
index 3270f6f785..4e7b78f628 100644
--- a/wrappers/python/openmm/app/simulation.py
+++ b/wrappers/python/openmm/app/simulation.py
@@ -206,9 +206,23 @@ def _simulate(self, endStep=None, endTime=None):
             
             anyReport = False
             for i, reporter in enumerate(self.reporters):
-                nextReport[i] = reporter.describeNextReport(self)
-                if nextReport[i][0] > 0 and nextReport[i][0] <= nextSteps:
-                    nextSteps = nextReport[i][0]
+                report = reporter.describeNextReport(self)
+                # convert to new dict format if the report is in the old tuple format
+                if isinstance(report, tuple):
+                    report_dict = {'steps': report[0]}
+
+                    if len(report) > 5:
+                        report_dict['periodic'] = report[5]
+
+                    includes = ['positions', 'velocities', 'forces', 'energy']
+                    report_dict['include'] = [includes[i] for i in range(4) if report[i+1]]
+                    report = report_dict
+                nextReport[i] = report
+                
+                steps = nextReport[i]['steps']
+
+                if steps > 0 and steps <= nextSteps:
+                    nextSteps = steps
                     anyReport = True
             stepsToGo = nextSteps
             while stepsToGo > 10:
@@ -226,19 +240,16 @@ def _simulate(self, endStep=None, endTime=None):
                 unwrapped = []
                 either = []
                 for reporter, report in zip(self.reporters, nextReport):
-                    if report[0] == nextSteps:
-                        if len(report) > 5:
-                            wantWrap = report[5]
-                            if wantWrap is None:
-                                wantWrap = self._usesPBC
-                        else:
-                            wantWrap = self._usesPBC
-                        if not report[1]:
+                    if report['steps'] == nextSteps:
+                        wantWrap = self._usesPBC if report.get('periodic') is None else report['periodic']
+
+                        if not 'positions' in report['include']: # if no positions are requested, we don't care about pbc
                             either.append((reporter, report))
                         elif wantWrap:
                             wrapped.append((reporter, report))
                         else:
                             unwrapped.append((reporter, report))
+
                 if len(wrapped) > len(unwrapped):
                     wrapped += either
                 else:
@@ -252,23 +263,22 @@ def _simulate(self, endStep=None, endTime=None):
                     self._generate_reports(unwrapped, False)
     
     def _generate_reports(self, reports, periodic):
-        getPositions = False
-        getVelocities = False
-        getForces = False
-        getEnergy = False
-        for reporter, next in reports:
-            if next[1]:
-                getPositions = True
-            if next[2]:
-                getVelocities = True
-            if next[3]:
-                getForces = True
-            if next[4]:
-                getEnergy = True
-        state = self.context.getState(positions=getPositions, velocities=getVelocities, forces=getForces,
-                                      energy=getEnergy, parameters=True, enforcePeriodicBox=periodic,
-                                      groups=self.context.getIntegrator().getIntegrationForceGroups())
-        for reporter, next in reports:
+        '''Generate reports for all requested reporters
+
+        Parameters
+        ----------
+        reports :  list of tuples
+                each tuple in the list contains a reporter object and a description of the next report produced by reporter.describeNextReport().
+                Note that the old format of returning a tuple from reporter.describeNextReport() is not supported in this function.
+        periodic : bool
+                Specifies whether particle positions should be translated so the center of every molecule lies in the same periodic box.
+        '''
+        
+        includes = set.union(*[set(report[1]['include']) for report in reports])
+        includeArgs = {property:True for property in includes}
+
+        state = self.context.getState(groups=self.context.getIntegrator().getIntegrationForceGroups(), enforcePeriodicBox=periodic, parameters=True, **includeArgs)
+        for reporter, nextReport in reports:
             reporter.report(self, state)
 
     def saveCheckpoint(self, file):
diff --git a/wrappers/python/openmm/app/statedatareporter.py b/wrappers/python/openmm/app/statedatareporter.py
index 64b3b14a79..9ea620bb23 100644
--- a/wrappers/python/openmm/app/statedatareporter.py
+++ b/wrappers/python/openmm/app/statedatareporter.py
@@ -158,6 +158,7 @@ def __init__(self, file, reportInterval, step=False, time=False, potentialEnergy
         self._needsVelocities = False
         self._needsForces = False
         self._needEnergy = potentialEnergy or kineticEnergy or totalEnergy or temperature
+        self._includes = ['energy'] if self._needEnergy else []
 
     def describeNextReport(self, simulation):
         """Get information about the next report this object will generate.
@@ -169,14 +170,11 @@ def describeNextReport(self, simulation):
 
         Returns
         -------
-        tuple
-            A five element tuple. The first element is the number of steps
-            until the next report. The remaining elements specify whether
-            that report will require positions, velocities, forces, and
-            energies respectively.
+        dict
+            A dictionary describing the required information for the next report
         """
         steps = self._reportInterval - simulation.currentStep%self._reportInterval
-        return (steps, self._needsPositions, self._needsVelocities, self._needsForces, self._needEnergy)
+        return {'steps':steps, 'periodic':None, 'include':self._includes}
 
     def report(self, simulation, state):
         """Generate a report.
diff --git a/wrappers/python/openmm/app/xtcreporter.py b/wrappers/python/openmm/app/xtcreporter.py
index c37d15e08a..38ef3b4d2d 100644
--- a/wrappers/python/openmm/app/xtcreporter.py
+++ b/wrappers/python/openmm/app/xtcreporter.py
@@ -42,15 +42,11 @@ def describeNextReport(self, simulation):
 
         Returns
         -------
-        tuple
-            A six element tuple. The first element is the number of steps
-            until the next report. The next four elements specify whether
-            that report will require positions, velocities, forces, and
-            energies respectively.  The final element specifies whether
-            positions should be wrapped to lie in a single periodic box.
+        dict
+            A dictionary describing the required information for the next report
         """
-        steps = self._reportInterval - simulation.currentStep % self._reportInterval
-        return (steps, True, False, False, False, self._enforcePeriodicBox)
+        steps = self._reportInterval - simulation.currentStep%self._reportInterval
+        return {'steps':steps, 'periodic':self._enforcePeriodicBox, 'include':['positions']}
 
     def report(self, simulation, state):
         """Generate a report.


Test patch:
diff --git a/wrappers/python/tests/TestSimulation.py b/wrappers/python/tests/TestSimulation.py
index aa26895634..68ba979670 100644
--- a/wrappers/python/tests/TestSimulation.py
+++ b/wrappers/python/tests/TestSimulation.py
@@ -179,7 +179,7 @@ def __init__(self, periodic):
                 
             def describeNextReport(self, simulation):
                 steps = self.interval - simulation.currentStep%self.interval
-                return (steps, True, False, False, False, self.periodic)
+                return {'steps':steps, 'periodic':self.periodic, 'include':['positions']}
         
             def report(self, simulation, state):
                 state2 = simulation.context.getState(getPositions=True, enforcePeriodicBox=self.periodic)