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

36.36

/mdsuite/calculators/einstein_helfand_thermal_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 thermal 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


@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 EinsteinHelfandThermalConductivity(TrajectoryCalculator, ABC):
    """
    Class for the Einstein-Helfand Ionic Conductivity.

    Attributes
    ----------
    experiment :  object
            Experiment class to call from
    x_label : str
            X label of the tensor_values when plotted
    y_label : str
            Y label of the tensor_values when plotted
    analysis_name : str
            Name of the analysis
    loaded_property : str
            Property loaded from the database_path for the analysis

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

    Examples
    --------
    experiment.run.EinsteinHelfandTThermalConductivity(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.integrated_heat_current
        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 Thermal Conductivity"
        self._dtype = tf.float64

        self.prefactor = None

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

        Parameters
        ----------
        plot : bool
                if true, plot the output.
        data_range : int
                Data range to use in the analysis.
        correlation_time : int
                Correlation time to use in the window sampling.
        """
        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 = 1
        denominator = (
            self.experiment.volume
            * self.experiment.temperature
            * self.experiment.units.boltzmann
        )
        units_change = (
            self.experiment.units.energy
            / self.experiment.units.length
            / self.experiment.units.time
            / self.experiment.units.temperature
        )
        self.prefactor = numerator / denominator * units_change

    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):
        """
        Calculate and return the msd.

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

        Returns
        -------
        MSD of the tensor_values.
        """
        msd = tf.math.squared_difference(ensemble, ensemble[None, 0])

        msd = self.prefactor * tf.reduce_sum(msd, axis=1)
        self.msd_array += np.array(msd)  # Update the averaged function

    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 = {
            "thermal_conductivity": 1 / 6 * fit_values[0],
            "uncertainty": 1 / 6 * error,
            "time": self.time.tolist(),
            "msd": self.msd_array.tolist(),
        }
        self.queue_data(data=data, subjects=["System"])

        # Update the plot if required
        if self.plot:
            self.run_visualization(
                x_data=np.array(self.time) * self.experiment.units.time,
                y_data=self.msd_array * self.experiment.units.time,
                title=f"{fit_values[0]} += {error}",
            )

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