Example M8: Compute trajectory average of a global attribute

This Python modifier lets the user specify the name of a global attribute for which to compute the time average over all frames of the current simulation trajectory. It provides a similar functionality as the built-in TimeAveragingModifier of OVITO Pro.

The modifier is based on the advanced programming interface, i.e., it is implemented as a Python class inheriting from ModifierInterface.

from ovito.pipeline import ModifierInterface
from ovito.data import DataCollection
from traits.api import Str

class TimeAverageOfAttribute(ModifierInterface):

    # Specifies the name of the global attribute for which to compute the trajectory time average.
    input_attribute = Str(default_value='', label='Input attribute')

    def modify(self, data: DataCollection, *, input_slots: dict[str, ModifierInterface.InputSlot], data_cache: DataCollection, **kwargs):

        # Display an info string in the status bar of OVITO Pro to indicate what we are doing.
        yield f'Computing time average of attribute {self.input_attribute}'

        if self.input_attribute == '':
            print("Please specify the name of the global attribute to be averaged.")
            return
        elif self.input_attribute not in data.attributes:
            print(f"Global attribute '{self.input_attribute}' does not exist.")
            return

        # Determine the name of the new output attribute, which will store the computed time average.
        output_attribute = self.input_attribute + ' (Average)'

        # Check whether the computed time average is already stored in the modifier's cache. If not, compute it now:
        if output_attribute not in data_cache.attributes:
            slot = input_slots['upstream']

            # Iterate over all frames of the input trajectory to sample instantaneous values of input attribute:
            sum = 0.0
            for frame in range(slot.num_frames):
                frame_data = slot.compute(frame)
                sum += frame_data.attributes[self.input_attribute]

                # Progress reporting: Percentage value will be displayed in the status bar of OVITO Pro.
                yield frame / slot.num_frames

            # Compute mean value and store it in the cache.
            data_cache.attributes[output_attribute] = sum / slot.num_frames

        # Copy value over from the cache to the actual output data collection:
        data.attributes[output_attribute] = data_cache.attributes[output_attribute]