Source code for mdsuite.time_series.base

"""
MDSuite: A Zincwarecode package.

License
-------
This program and the accompanying materials are made available under the terms
of the Eclipse Public License v2.0 which accompanies this distribution, and is
available at https://www.eclipse.org/legal/epl-v20.html

SPDX-License-Identifier: EPL-2.0

Copyright Contributors to the Zincwarecode Project.

Contact Information
-------------------
email: zincwarecode@gmail.com
github: https://github.com/zincware
web: https://zincwarecode.com/

Citation
--------
If you use this module please cite us with:

Summary
-------
"""
from __future__ import annotations

from typing import TYPE_CHECKING

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

from mdsuite.database.simulation_database import Database

if TYPE_CHECKING:
    from mdsuite import Experiment


[docs]def running_mean(x, N): """Perform a rolling window mean""" cumsum = np.cumsum(np.insert(x, 0, 0)) return (cumsum[N:] - cumsum[:-N]) / float(N)
[docs]class TimeSeries: def __init__(self, experiment: Experiment): """ Parameters ---------- experiment: Experiment The parent experiment class to perform the time series operation on """ self.experiment = experiment self.loaded_property = None self.fig_labels = {"x": None, "y": None} self.species = experiment.species self.rolling_window = 0 self.reduce_sum = True # Properties self._database = None self._data = None def __call__(self, species: list = None, rolling_window: int = 0): if species is not None: self.species = species self.rolling_window = rolling_window self.plot() @property def database(self): """Get the database""" if self._database is None: self._database = Database(self.experiment.database_path / "database.hdf5") return self._database @property def data(self): """Get the data for all species and timesteps for the loaded_property""" if self._data is None: self._data = tf.concat( [ self.database.load_data([f"{species}/{self.loaded_property}"]) for species in self.species ], axis=0, ) return self._data @property def preprocess_data(self): """Perform some data preprocessing before plotting it""" data = self.data if self.reduce_sum: data = tf.einsum("atx -> t", data) # perform a reduce sum over atoms "a" and property dimension "x" to # yield time steps "t" if self.rolling_window > 0: data = running_mean(data, self.rolling_window) return data
[docs] def plot(self): """Plot the data over timesteps""" fig, ax = plt.subplots() ax.plot(self.preprocess_data) ax.set_xlabel(self.fig_labels["x"]) ax.set_ylabel(self.fig_labels["y"]) fig.show()