3578337630

Build Type

github-actions

Committed by GitHub

Commit Message

<a href="https://github.com/zincware/MDSuite/commit/<a class=hub.com/zincware/MDSuite/commit/<a class="double-link" href="https://git"><a class=hub.com/zincware/MDSuite/commit/<a class="double-link" href="https://git"><a class=hub.com/zincware/MDSuite/commit/<a class="double-link" href="https://git"><a class=hub.com/zincware/MDSuite/commit/<a class="double-link" href="https://git"><a class=hub.com/zincware/MDSuite/commit/7c693ed5f3d0d1a79f4c687a3befd246b495316c">7c693ed5f">&lt;a href=&quot;https://github.com/zincware/MDSuite/commit/</a><a class="double-link" href="https://github.com/zincware/MDSuite/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/zincware/MDSuite/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/zincware/MDSuite/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/zincware/MDSuite/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/zincware/MDSuite/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/zincware/MDSuite/commit/7c693ed5f3d0d1a79f4c687a3befd246b495316c">7c693ed5f</a><a href="https://github.com/zincware/MDSuite/commit/7c693ed5f3d0d1a79f4c687a3befd246b495316c">&quot;&gt;&amp;lt;a href=&amp;quot;https://github.com/zincware/MDSuite/commit/&lt;/a&gt;&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/zincware/MDSuite/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/zincware/MDSuite/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/zincware/MDSuite/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/zincware/MDSuite/commit/7c693ed5f3d0d1a79f4c687a3befd246b495316c&quot;&gt;7c693ed5f&lt;/a&gt;&lt;a href=&quot;https://github.com/zincware/MDSuite/commit/7c693ed5f3d0d1a79f4c687a3befd246b495316c&quot;&gt;&amp;lt;a href=&amp;quot;https://github.com/zincware/MDSuite/commit/7c693ed5f3d0d1a79f4c687a3befd246b495316c&amp;quot;&amp;gt;&amp;amp;quot;&amp;amp;gt;&amp;amp;amp;lt;a href=&amp;amp;amp;quot;https://github.com/zincware/MDSuite/commit/&amp;amp;lt;/a&amp;amp;gt;&amp;amp;lt;a class=&... (continued)

Pull Request Pull Request #564: [WIP] Sam tov fix memory scaling tests

Run Details

912 of 1250 branches covered (72.96%)

Branch coverage included in aggregate %.

13 of 13 new or added lines in 3 files covered. (100.0%)

3931 of 4823 relevant lines covered (81.51%)

2.44 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.67

/mdsuite/calculators/einstein_diffusion_coefficients.py

"""
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
-------
Module for the computation of self-diffusion coefficients using the Einstein method.
"""
from __future__ import annotations

import logging
from abc import ABC
from dataclasses import dataclass
from typing import Any, List, Union

import numpy as np
import tensorflow as tf
from bokeh.models import HoverTool, LinearAxis, Span
from bokeh.models.ranges import Range1d
from bokeh.plotting import figure
from tqdm import tqdm

from mdsuite import utils
from mdsuite.calculators.calculator import call
from mdsuite.calculators.trajectory_calculator import TrajectoryCalculator
from mdsuite.database.mdsuite_properties import mdsuite_properties
from mdsuite.utils.calculator_helper_methods import fit_einstein_curve

log = logging.getLogger(__name__)


@dataclass
class Args:
    """
    Data class for the saved properties.
    """

    data_range: int
    correlation_time: int
    atom_selection: np.s_
    tau_values: np.s_
    molecules: bool
    species: list
    fit_range: int


class EinsteinDiffusionCoefficients(TrajectoryCalculator, ABC):
    """
    Class for the Einstein diffusion coefficient implementation

    Attributes
    ----------
    msd_array : np.ndarray
            MSd data updated during each ensemble computation.

    See Also
    --------
    mdsuite.calculators.calculator.Calculator class

    Examples
    --------
    project.experiment.run.EinsteinDiffusionCoefficients(data_range=500,
                                                         plot=True,
                                                         correlation_time=10)
    """

    def __init__(self, **kwargs):
        """
        Parameters
        ----------
        experiment :  Experiment
                Experiment class to call from
        experiments :  Experiment
                Experiment classes to call from
        """

        super().__init__(**kwargs)
        self.scale_function = {"linear": {"scale_factor": 150}}
        self.loaded_property = mdsuite_properties.unwrapped_positions
        self.x_label = r"$$\text{Time} / s$$"
        self.y_label = r"$$\text{MSD} / m^{2}$$"
        self.result_keys = [
            "diffusion_coefficient",
            "uncertainty",
            "gradient",
            "intercept",
        ]
        self.result_series_keys = ["time", "msd", "gradients", "gradient_errors"]
        self.analysis_name = "Einstein Self-Diffusion Coefficients"
        self._dtype = tf.float64

        self.msd_array = None

        log.info("starting Einstein Diffusion Computation")

    @call
    def __call__(
        self,
        plot: bool = True,
        species: list = None,
        data_range: int = 100,
        correlation_time: int = 1,
        atom_selection: np.s_ = np.s_[:],
        molecules: bool = False,
        tau_values: Union[int, List, Any] = np.s_[:],
        fit_range: int = -1,
    ):
        """

        Parameters
        ----------
        plot : bool
                if true, plot the output.
        species : list
                List of species on which to operate.
        data_range : int
                Data range to use in the analysis.
        correlation_time : int
                Correlation time to use in the window sampling.
        atom_selection : np.s_
                Selection of atoms to use within the HDF5 database.
        molecules : bool
                If true, molecules are used instead of atoms.
        tau_values : Union[int, list, np.s_]
                Selection of tau values to use in the window sliding.

        Returns
        -------
        None
        """
        if species is None:
            if molecules:
                species = list(self.experiment.molecules)
            else:
                species = list(self.experiment.species)

        if fit_range == -1:
            fit_range = int(data_range - 1)
        # set args that will affect the computation result
        self.args = Args(
            data_range=data_range,
            correlation_time=correlation_time,
            atom_selection=atom_selection,
            tau_values=tau_values,
            molecules=molecules,
            species=species,
            fit_range=fit_range,
        )
        self.plot = plot
        self.system_property = False

    def ensemble_operation(self, ensemble):
        """
        Calculate and return the msd.

        Parameters
        ----------
        ensemble : tf.Tensor
                An ensemble of data to be operated on.

        Returns
        -------
        MSD of the tensor_values.
        """
        msd = tf.math.squared_difference(
            tf.gather(ensemble, self.args.tau_values, axis=1), ensemble[:, None, 0]
        )
        self.count += msd.shape[0]
        # average over particles, sum over dimensions
        # msd = tf.reduce_sum(tf.reduce_mean(msd, axis=0), axis=-1)
        msd = tf.reduce_sum(tf.reduce_sum(msd, axis=0), axis=-1)

        # sum up ensembles to average in post processing
        return np.array(msd)

    def fit_diff_coeff(self):
        """
        Apply unit conversion, fit line to the data, prepare for database storage
        """

        # self.msd_array /= int(self.n_batches) * self.ensemble_loop
        self.msd_array /= self.count

        self.msd_array *= self.experiment.units.length**2
        self.time *= self.experiment.units.time

        fit_values, covariance, gradients, gradient_errors = fit_einstein_curve(
            x_data=self.time, y_data=self.msd_array, fit_max_index=self.args.fit_range
        )
        error = np.sqrt(np.diag(covariance))[0]

        data = {
            self.result_keys[0]: 1 / 6.0 * fit_values[0],
            self.result_keys[1]: 1 / 6.0 * error,
            self.result_keys[2]: fit_values[0],
            self.result_keys[3]: fit_values[1],
            self.result_series_keys[0]: self.time.tolist(),
            self.result_series_keys[1]: self.msd_array.tolist(),
            self.result_series_keys[2]: (np.array(gradients) / 6).tolist(),
            self.result_series_keys[3]: (np.array(gradient_errors) / 6).tolist(),
        }
        return data

    def run_calculator(self):
        """
        Run analysis.
        """
        self._run_dependency_check()
        for species in self.args.species:
            # Here for now to avoid issues. Should be moved out when calculators become
            # species-wise
            self.time = None
            self.time = self._handle_tau_values()
            dict_ref = str.encode("/".join([species, self.loaded_property.name]))
            batch_ds = self.get_batch_dataset([species])
            self.msd_array = np.zeros(self.data_resolution)
            self.count = 0
            # loop over batches to get MSD
            for i, batch in tqdm(
                enumerate(batch_ds),
                ncols=70,
                desc=species,
                total=self.n_batches,
                disable=self.memory_manager.minibatch,
            ):
                ensemble_ds = self.get_ensemble_dataset(batch, species)

                for ensemble in ensemble_ds:
                    if not ensemble[dict_ref].shape[1] == self.args.data_range:
                        continue
                    else:
                        self.msd_array += self.ensemble_operation(ensemble[dict_ref])
                        self.count += 1

            # self.msd_array = np.array(tf.reduce_sum(self.msd_array, axis=0))
            fit_results = self.fit_diff_coeff()
            self.queue_data(data=fit_results, subjects=[species])

    def plot_data(self, data):
        """
        Plot the Einstein fits.

        Parameters
        ----------
        data

        Returns
        -------

        """
        for selected_species, val in data.items():
            fig = figure(x_axis_label=self.x_label, y_axis_label=self.y_label)

            gradients = np.array(val[self.result_series_keys[2]])
            gradient_errors = np.array(val[self.result_series_keys[3]])

            time = np.array(val[self.result_series_keys[0]])
            msd = np.array(val[self.result_series_keys[1]])

            fig.y_range = Range1d(-0.0, 1.1 * max(msd))

            # Compute the span
            span = Span(
                location=time[self.args.fit_range],
                dimension="height",
                line_dash="dashed",
            )
            # Compute msd and fit lines
            fig.line(
                time,
                msd,
                color=utils.Colour.ORANGE,
                legend_label=(
                    f"{selected_species}: {val[self.result_keys[0]]: 0.3E} +-"
                    f" {val[self.result_keys[1]]: 0.3E}"
                ),
            )
            fit_data = val[self.result_keys[2]] * time + val[self.result_keys[3]]
            fig.line(time, fit_data, color=utils.Colour.PAUA, legend_label="Curve fit.")
            fig.extra_y_ranges = {
                "Diff_range": Range1d(
                    start=0.999 * min(gradients), end=1.001 * max(gradients)
                )
            }

            fig.add_layout(
                LinearAxis(
                    y_range_name="Diff_range",
                    axis_label=r"$$\text{Diffusion Coefficient} / m^{2}s^{-1}$$",
                ),
                "right",
            )
            grad_time = time[-len(gradients) :]
            fig.line(
                grad_time,
                gradients,
                y_range_name="Diff_range",
                color=utils.Colour.MULBERRY,
            )
            fig.varea(
                grad_time,
                gradients - gradient_errors,
                gradients + gradient_errors,
                alpha=0.3,
                color=utils.Colour.ORANGE,
                y_range_name="Diff_range",
            )

            fig.add_tools(HoverTool())
            fig.add_layout(span)
            self.plot_array.append(fig)

1	"""
2	MDSuite: A Zincwarecode package.
3
4	License
5	-------
6	This program and the accompanying materials are made available under the terms
7	of the Eclipse Public License v2.0 which accompanies this distribution, and is
8	available at https://www.eclipse.org/legal/epl-v20.html
9
10	SPDX-License-Identifier: EPL-2.0
11
12	Copyright Contributors to the Zincwarecode Project.
13
14	Contact Information
15	-------------------
16	email: zincwarecode@gmail.com
17	github: https://github.com/zincware
18	web: https://zincwarecode.com/
19
20	Citation
21	--------
22	If you use this module please cite us with:
23
24	Summary
25	-------
26	Module for the computation of self-diffusion coefficients using the Einstein method.
27	"""
28	from __future__ import annotations	3✔
29
30	import logging	3✔
31	from abc import ABC	3✔
32	from dataclasses import dataclass	3✔
33	from typing import Any, List, Union	3✔
34
35	import numpy as np	3✔
36	import tensorflow as tf	3✔
37	from bokeh.models import HoverTool, LinearAxis, Span	3✔
38	from bokeh.models.ranges import Range1d	3✔
39	from bokeh.plotting import figure	3✔
40	from tqdm import tqdm	3✔
41
42	from mdsuite import utils	3✔
43	from mdsuite.calculators.calculator import call	3✔
44	from mdsuite.calculators.trajectory_calculator import TrajectoryCalculator	3✔
45	from mdsuite.database.mdsuite_properties import mdsuite_properties	3✔
46	from mdsuite.utils.calculator_helper_methods import fit_einstein_curve	3✔
47
48	log = logging.getLogger(__name__)	3✔
49
50
51	@dataclass	3✔
52	class Args:	2✔
53	"""
54	Data class for the saved properties.
55	"""
56
57	data_range: int	3✔
58	correlation_time: int	3✔
59	atom_selection: np.s_	3✔
60	tau_values: np.s_	3✔
61	molecules: bool	3✔
62	species: list	3✔
63	fit_range: int	3✔
64
65
66	class EinsteinDiffusionCoefficients(TrajectoryCalculator, ABC):	3✔
67	"""
68	Class for the Einstein diffusion coefficient implementation
69
70	Attributes
71	----------
72	msd_array : np.ndarray
73	MSd data updated during each ensemble computation.
74
75	See Also
76	--------
77	mdsuite.calculators.calculator.Calculator class
78
79	Examples
80	--------
81	project.experiment.run.EinsteinDiffusionCoefficients(data_range=500,
82	plot=True,
83	correlation_time=10)
84	"""
85
86	def __init__(self, **kwargs):	3✔
87	"""
88	Parameters
89	----------
90	experiment : Experiment
91	Experiment class to call from
92	experiments : Experiment
93	Experiment classes to call from
94	"""
95
96	super().__init__(**kwargs)	3✔
97	self.scale_function = {"linear": {"scale_factor": 150}}	3✔
98	self.loaded_property = mdsuite_properties.unwrapped_positions	3✔
99	self.x_label = r"$$\text{Time} / s$$"	3✔
100	self.y_label = r"$$\text{MSD} / m^{2}$$"	3✔
101	self.result_keys = [	3✔
102	"diffusion_coefficient",
103	"uncertainty",
104	"gradient",
105	"intercept",
106	]
107	self.result_series_keys = ["time", "msd", "gradients", "gradient_errors"]	3✔
108	self.analysis_name = "Einstein Self-Diffusion Coefficients"	3✔
109	self._dtype = tf.float64	3✔
110
111	self.msd_array = None	3✔
112
113	log.info("starting Einstein Diffusion Computation")	3✔
114
115	@call	3✔
116	def __call__(	3✔
117	self,
118	plot: bool = True,
119	species: list = None,
120	data_range: int = 100,
121	correlation_time: int = 1,
122	atom_selection: np.s_ = np.s_[:],
123	molecules: bool = False,
124	tau_values: Union[int, List, Any] = np.s_[:],
125	fit_range: int = -1,
126	):
127	"""
128
129	Parameters
130	----------
131	plot : bool
132	if true, plot the output.
133	species : list
134	List of species on which to operate.
135	data_range : int
136	Data range to use in the analysis.
137	correlation_time : int
138	Correlation time to use in the window sampling.
139	atom_selection : np.s_
140	Selection of atoms to use within the HDF5 database.
141	molecules : bool
142	If true, molecules are used instead of atoms.
143	tau_values : Union[int, list, np.s_]
144	Selection of tau values to use in the window sliding.
145
146	Returns
147	-------
148	None
149	"""
150	if species is None:	3!
151	if molecules:	3✔
152	species = list(self.experiment.molecules)	3✔
153	else:
154	species = list(self.experiment.species)	3✔
155
156	if fit_range == -1:	3!
157	fit_range = int(data_range - 1)	3✔
158	# set args that will affect the computation result
159	self.args = Args(	3✔
160	data_range=data_range,
161	correlation_time=correlation_time,
162	atom_selection=atom_selection,
163	tau_values=tau_values,
164	molecules=molecules,
165	species=species,
166	fit_range=fit_range,
167	)
168	self.plot = plot	3✔
169	self.system_property = False	3✔
170
171	def ensemble_operation(self, ensemble):	3✔
172	"""
173	Calculate and return the msd.
174
175	Parameters
176	----------
177	ensemble : tf.Tensor
178	An ensemble of data to be operated on.
179
180	Returns
181	-------
182	MSD of the tensor_values.
183	"""
184	msd = tf.math.squared_difference(	3✔
185	tf.gather(ensemble, self.args.tau_values, axis=1), ensemble[:, None, 0]
186	)
187	self.count += msd.shape[0]	3✔
188	# average over particles, sum over dimensions
189	# msd = tf.reduce_sum(tf.reduce_mean(msd, axis=0), axis=-1)
190	msd = tf.reduce_sum(tf.reduce_sum(msd, axis=0), axis=-1)	3✔
191
192	# sum up ensembles to average in post processing
193	return np.array(msd)	3✔
194
195	def fit_diff_coeff(self):	3✔
196	"""
197	Apply unit conversion, fit line to the data, prepare for database storage
198	"""
199
200	# self.msd_array /= int(self.n_batches) * self.ensemble_loop
201	self.msd_array /= self.count	3✔
202
203	self.msd_array = self.experiment.units.length*2	3✔
204	self.time *= self.experiment.units.time	3✔
205
206	fit_values, covariance, gradients, gradient_errors = fit_einstein_curve(	3✔
207	x_data=self.time, y_data=self.msd_array, fit_max_index=self.args.fit_range
208	)
209	error = np.sqrt(np.diag(covariance))[0]	3✔
210
211	data = {	3✔
212	self.result_keys[0]: 1 / 6.0 * fit_values[0],
213	self.result_keys[1]: 1 / 6.0 * error,
214	self.result_keys[2]: fit_values[0],
215	self.result_keys[3]: fit_values[1],
216	self.result_series_keys[0]: self.time.tolist(),
217	self.result_series_keys[1]: self.msd_array.tolist(),
218	self.result_series_keys[2]: (np.array(gradients) / 6).tolist(),
219	self.result_series_keys[3]: (np.array(gradient_errors) / 6).tolist(),
220	}
221	return data	3✔
222
223	def run_calculator(self):	3✔
224	"""
225	Run analysis.
226	"""
227	self._run_dependency_check()	3✔
228	for species in self.args.species:	3✔
229	# Here for now to avoid issues. Should be moved out when calculators become
230	# species-wise
231	self.time = None	3✔
232	self.time = self._handle_tau_values()	3✔
233	dict_ref = str.encode("/".join([species, self.loaded_property.name]))	3✔
234	batch_ds = self.get_batch_dataset([species])	3✔
235	self.msd_array = np.zeros(self.data_resolution)	3✔
236	self.count = 0	3✔
237	# loop over batches to get MSD
238	for i, batch in tqdm(	3✔
239	enumerate(batch_ds),
240	ncols=70,
241	desc=species,
242	total=self.n_batches,
243	disable=self.memory_manager.minibatch,
244	):
245	ensemble_ds = self.get_ensemble_dataset(batch, species)	3✔
246
247	for ensemble in ensemble_ds:	3✔
248	if not ensemble[dict_ref].shape[1] == self.args.data_range:	3!
249	continue	×
250	else:
251	self.msd_array += self.ensemble_operation(ensemble[dict_ref])	3✔
252	self.count += 1	3✔
253
254	# self.msd_array = np.array(tf.reduce_sum(self.msd_array, axis=0))
255	fit_results = self.fit_diff_coeff()	3✔
256	self.queue_data(data=fit_results, subjects=[species])	3✔
257
258	def plot_data(self, data):	3✔
259	"""
260	Plot the Einstein fits.
261
262	Parameters
263	----------
264	data
265
266	Returns
267	-------
268
269	"""
270	for selected_species, val in data.items():	3✔
271	fig = figure(x_axis_label=self.x_label, y_axis_label=self.y_label)	3✔
272
273	gradients = np.array(val[self.result_series_keys[2]])	3✔
274	gradient_errors = np.array(val[self.result_series_keys[3]])	3✔
275
276	time = np.array(val[self.result_series_keys[0]])	3✔
277	msd = np.array(val[self.result_series_keys[1]])	3✔
278
279	fig.y_range = Range1d(-0.0, 1.1 * max(msd))	3✔
280
281	# Compute the span
282	span = Span(	3✔
283	location=time[self.args.fit_range],
284	dimension="height",
285	line_dash="dashed",
286	)
287	# Compute msd and fit lines
288	fig.line(	3✔
289	time,
290	msd,
291	color=utils.Colour.ORANGE,
292	legend_label=(
293	f"{selected_species}: {val[self.result_keys[0]]: 0.3E} +-"
294	f" {val[self.result_keys[1]]: 0.3E}"
295	),
296	)
297	fit_data = val[self.result_keys[2]] * time + val[self.result_keys[3]]	3✔
298	fig.line(time, fit_data, color=utils.Colour.PAUA, legend_label="Curve fit.")	3✔
299	fig.extra_y_ranges = {	3✔
300	"Diff_range": Range1d(
301	start=0.999 * min(gradients), end=1.001 * max(gradients)
302	)
303	}
304
305	fig.add_layout(	3✔
306	LinearAxis(
307	y_range_name="Diff_range",
308	axis_label=r"$$\text{Diffusion Coefficient} / m^{2}s^{-1}$$",
309	),
310	"right",
311	)
312	grad_time = time[-len(gradients) :]	3✔
313	fig.line(	3✔
314	grad_time,
315	gradients,
316	y_range_name="Diff_range",
317	color=utils.Colour.MULBERRY,
318	)
319	fig.varea(	3✔
320	grad_time,
321	gradients - gradient_errors,
322	gradients + gradient_errors,
323	alpha=0.3,
324	color=utils.Colour.ORANGE,
325	y_range_name="Diff_range",
326	)
327
328	fig.add_tools(HoverTool())	3✔
329	fig.add_layout(span)	3✔
330	self.plot_array.append(fig)	3✔

zincware / MDSuite / 3578337630

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous