3999396905

Build Type

push

github-actions

Committed by GitHub

Commit Message

[merge before other PRs] ruff updates (#580)

Run Details

960 of 1311 branches covered (73.23%)

Branch coverage included in aggregate %.

15 of 15 new or added lines in 11 files covered. (100.0%)

4034 of 4930 relevant lines covered (81.83%)

3.19 hits per line

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

39.19

/mdsuite/calculators/einstein_helfand_ionic_conductivity.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
-------
MDSuite module for the computation of ionic conductivity using the Einstein method.
"""
from abc import ABC
from dataclasses import dataclass

import numpy as np
import tensorflow as tf
from tqdm import tqdm

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
from mdsuite.utils.units import boltzmann_constant, elementary_charge


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

    data_range: int
    correlation_time: int
    tau_values: np.s_
    atom_selection: np.s_
    fit_range: int


class EinsteinHelfandIonicConductivity(TrajectoryCalculator, ABC):
    """
    Class for the Einstein-Helfand Ionic Conductivity.

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

    Examples
    --------
    experiment.run_computation.EinsteinHelfandTIonicConductivity(data_range=500,
                                                                 plot=True,
                                                                 correlation_time=10)
    """

    def __init__(self, **kwargs):
        """
        Python constructor.

        Parameters
        ----------
        experiment :  object
            Experiment class to call from
        """
        # parse to the experiment class
        super().__init__(**kwargs)
        self.scale_function = {"linear": {"scale_factor": 5}}

        self.loaded_property = mdsuite_properties.translational_dipole_moment
        self.dependency = mdsuite_properties.unwrapped_positions
        self.system_property = True

        self.x_label = r"$$\text{Time} / s$$"
        self.y_label = r"$$\text{MSD} / m^2/s$$"
        self.analysis_name = "Einstein Helfand Ionic Conductivity"

        self.result_keys = ["ionic_conductivity", "uncertainty"]
        self.result_series_keys = ["time", "msd"]

        self._dtype = tf.float64

    @call
    def __call__(
        self,
        plot=True,
        data_range=100,
        correlation_time=1,
        tau_values: np.s_ = np.s_[:],
        fit_range: int = -1,
    ):
        """
        Python constructor.

        Parameters
        ----------
        plot : bool
                if true, plot the tensor_values
        data_range :
                Number of configurations to use in each ensemble
        correlation_time : int
                Correlation time to use in the analysis.
        """
        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,
            tau_values=tau_values,
            atom_selection=np.s_[:],
            fit_range=fit_range,
        )

        self.plot = plot
        self.time = self._handle_tau_values()
        self.msd_array = np.zeros(self.data_resolution)

    def check_input(self):
        """
        Check the user input to ensure no conflicts are present.

        Returns
        -------

        """
        self._run_dependency_check()

    def _calculate_prefactor(self):
        """
        Compute the ionic conductivity prefactor.

        Returns
        -------

        """
        # Calculate the prefactor
        numerator = (self.experiment.units.length**2) * (elementary_charge**2)
        denominator = (
            self.experiment.units.time
            * self.experiment.volume
            * self.experiment.units.volume
            * self.experiment.temperature
            * boltzmann_constant
        )
        self.prefactor = numerator / denominator

    def _apply_averaging_factor(self):
        """
        Apply the averaging factor to the msd array.

        Returns
        -------
        -------.

        """
        self.msd_array /= int(self.n_batches) * self.ensemble_loop

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

        Parameters
        ----------
        ensemble

        Returns
        -------
        MSD of the tensor_values.
        """
        msd = tf.math.squared_difference(
            tf.gather(ensemble, self.args.tau_values, axis=1), ensemble[:, 0, :]
        )
        msd = self.prefactor * tf.reduce_sum(msd, axis=2)
        self.msd_array += np.array(msd)[0, :]

    def _post_operation_processes(self):
        """
        call the post-op processes
        Returns
        -------.

        """
        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 * fit_values[0],
            self.result_keys[1]: 1 / 6 * error,
            self.result_series_keys[0]: self.time.tolist(),
            self.result_series_keys[1]: self.msd_array.tolist(),
        }

        self.queue_data(data=data, subjects=["System"])

    def run_calculator(self):
        """

        Run analysis.

        Returns
        -------

        """
        self.check_input()
        # Compute the pre-factor early.
        self._calculate_prefactor()

        dict_ref = str.encode(
            "/".join([self.loaded_property.name, self.loaded_property.name])
        )

        batch_ds = self.get_batch_dataset([self.loaded_property.name])

        for batch in tqdm(
            batch_ds,
            ncols=70,
            total=self.n_batches,
            disable=self.memory_manager.minibatch,
        ):
            ensemble_ds = self.get_ensemble_dataset(batch, self.loaded_property.name)

            for ensemble in ensemble_ds:
                self.ensemble_operation(ensemble[dict_ref])

        # Scale, save, and plot the data.
        self._apply_averaging_factor()
        self._post_operation_processes()

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	MDSuite module for the computation of ionic conductivity using the Einstein method.
27	"""
28	from abc import ABC	4✔
29	from dataclasses import dataclass	4✔
30
31	import numpy as np	4✔
32	import tensorflow as tf	4✔
33	from tqdm import tqdm	4✔
34
35	from mdsuite.calculators.calculator import call	4✔
36	from mdsuite.calculators.trajectory_calculator import TrajectoryCalculator	4✔
37	from mdsuite.database.mdsuite_properties import mdsuite_properties	4✔
38	from mdsuite.utils.calculator_helper_methods import fit_einstein_curve	4✔
39	from mdsuite.utils.units import boltzmann_constant, elementary_charge	4✔
40
41
42	@dataclass	4✔
43	class Args:	3✔
44	"""Data class for the saved properties."""
45
46	data_range: int	4✔
47	correlation_time: int	4✔
48	tau_values: np.s_	4✔
49	atom_selection: np.s_	4✔
50	fit_range: int	4✔
51
52
53	class EinsteinHelfandIonicConductivity(TrajectoryCalculator, ABC):	4✔
54	"""
55	Class for the Einstein-Helfand Ionic Conductivity.
56
57	See Also
58	--------
59	mdsuite.calculators.calculator.Calculator class
60
61	Examples
62	--------
63	experiment.run_computation.EinsteinHelfandTIonicConductivity(data_range=500,
64	plot=True,
65	correlation_time=10)
66	"""
67
68	def __init__(self, **kwargs):	4✔
69	"""
70	Python constructor.
71
72	Parameters
73	----------
74	experiment : object
75	Experiment class to call from
76	"""
77	# parse to the experiment class
78	super().__init__(**kwargs)	×
79	self.scale_function = {"linear": {"scale_factor": 5}}	×
80
81	self.loaded_property = mdsuite_properties.translational_dipole_moment	×
82	self.dependency = mdsuite_properties.unwrapped_positions	×
83	self.system_property = True	×
84
85	self.x_label = r"$$\text{Time} / s$$"	×
86	self.y_label = r"$$\text{MSD} / m^2/s$$"	×
87	self.analysis_name = "Einstein Helfand Ionic Conductivity"	×
88
89	self.result_keys = ["ionic_conductivity", "uncertainty"]	×
90	self.result_series_keys = ["time", "msd"]	×
91
92	self._dtype = tf.float64	×
93
94	@call	4✔
95	def __call__(	4✔
96	self,
97	plot=True,
98	data_range=100,
99	correlation_time=1,
100	tau_values: np.s_ = np.s_[:],
101	fit_range: int = -1,
102	):
103	"""
104	Python constructor.
105
106	Parameters
107	----------
108	plot : bool
109	if true, plot the tensor_values
110	data_range :
111	Number of configurations to use in each ensemble
112	correlation_time : int
113	Correlation time to use in the analysis.
114	"""
115	if fit_range == -1:	×
116	fit_range = int(data_range - 1)	×
117
118	# set args that will affect the computation result
119	self.args = Args(	×
120	data_range=data_range,
121	correlation_time=correlation_time,
122	tau_values=tau_values,
123	atom_selection=np.s_[:],
124	fit_range=fit_range,
125	)
126
127	self.plot = plot	×
128	self.time = self._handle_tau_values()	×
129	self.msd_array = np.zeros(self.data_resolution)	×
130
131	def check_input(self):	4✔
132	"""
133	Check the user input to ensure no conflicts are present.
134
135	Returns
136	-------
137
138	"""
139	self._run_dependency_check()	×
140
141	def _calculate_prefactor(self):	4✔
142	"""
143	Compute the ionic conductivity prefactor.
144
145	Returns
146	-------
147
148	"""
149	# Calculate the prefactor
150	numerator = (self.experiment.units.length*2) (elementary_charge**2)	×
151	denominator = (	×
152	self.experiment.units.time
153	* self.experiment.volume
154	* self.experiment.units.volume
155	* self.experiment.temperature
156	* boltzmann_constant
157	)
158	self.prefactor = numerator / denominator	×
159
160	def _apply_averaging_factor(self):	4✔
161	"""
162	Apply the averaging factor to the msd array.
163
164	Returns
165	-------
166	-------.
167
168	"""
169	self.msd_array /= int(self.n_batches) * self.ensemble_loop	×
170
171	def ensemble_operation(self, ensemble: tf.Tensor):	4✔
172	"""
173	Calculate and return the msd.
174
175	Parameters
176	----------
177	ensemble
178
179	Returns
180	-------
181	MSD of the tensor_values.
182	"""
183	msd = tf.math.squared_difference(	×
184	tf.gather(ensemble, self.args.tau_values, axis=1), ensemble[:, 0, :]
185	)
186	msd = self.prefactor * tf.reduce_sum(msd, axis=2)	×
187	self.msd_array += np.array(msd)[0, :]	×
188
189	def _post_operation_processes(self):	4✔
190	"""
191	call the post-op processes
192	Returns
193	-------.
194
195	"""
196	fit_values, covariance, gradients, gradient_errors = fit_einstein_curve(	×
197	x_data=self.time, y_data=self.msd_array, fit_max_index=self.args.fit_range
198	)
199	error = np.sqrt(np.diag(covariance))[0]	×
200
201	data = {	×
202	self.result_keys[0]: 1 / 6 * fit_values[0],
203	self.result_keys[1]: 1 / 6 * error,
204	self.result_series_keys[0]: self.time.tolist(),
205	self.result_series_keys[1]: self.msd_array.tolist(),
206	}
207
208	self.queue_data(data=data, subjects=["System"])	×
209
210	def run_calculator(self):	4✔
211	"""
212
213	Run analysis.
214
215	Returns
216	-------
217
218	"""
219	self.check_input()	×
220	# Compute the pre-factor early.
221	self._calculate_prefactor()	×
222
223	dict_ref = str.encode(	×
224	"/".join([self.loaded_property.name, self.loaded_property.name])
225	)
226
227	batch_ds = self.get_batch_dataset([self.loaded_property.name])	×
228
229	for batch in tqdm(	×
230	batch_ds,
231	ncols=70,
232	total=self.n_batches,
233	disable=self.memory_manager.minibatch,
234	):
235	ensemble_ds = self.get_ensemble_dataset(batch, self.loaded_property.name)	×
236
237	for ensemble in ensemble_ds:	×
238	self.ensemble_operation(ensemble[dict_ref])	×
239
240	# Scale, save, and plot the data.
241	self._apply_averaging_factor()	×
242	self._post_operation_processes()	×

zincware / MDSuite / 3999396905

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