15077298886

import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
from matplotlib import axes, figure
from matplotlib.ticker import MaxNLocator

from dfttoolkit.output import AimsOutput


class VisualiseAims(AimsOutput):
    """
    FHI-aims visualisation tools.

    ...

    Attributes
    ----------
    scf_conv_acc_params : dict
        the SCF convergence accuracy parameters
    """

    @staticmethod
    def _plot_charge_convergence(
        ax: axes.Axes,
        tot_scf_iters: npt.NDArray[np.int64] | list[int],
        delta_charge: npt.NDArray[np.float64] | list[float],
        delta_charge_sd: npt.NDArray[np.float64] | list[float] | None = None,
        conv_params: dict | None = None,
        title: str | None = None,
    ) -> axes.Axes:
        """
        Create a subplot for the charge convergence of an FHI-aims calculation.

        Parameters
        ----------
        ax : axes.Axes
            matplotlib subplot index
        tot_scf_iters : npt.NDArray[np.int64] | list[int]
            cumulative SCF iterations
        delta_charge : npt.NDArray[np.float64] | list[float]
            change of spin-up or total spin (if the calculation was spin none)
            eigenvalues
        delta_charge_sd : npt.NDArray[np.float64] | list[float] | None, default=None
            change of spin-down eigenvalues
        conv_params : dict | None, default=None
            convergence parameters which determine if the SCF cycle has converged
        title : str | None, default=None
            system name to include in title

        Returns
        -------
        axes.Axes
            matplotlib subplot object
        """
        ax.plot(tot_scf_iters, delta_charge, label=r"$\Delta$ charge")

        # Only plot delta_charge_sd if the calculation is spin polarised
        if delta_charge_sd is not None:
            ax.plot(
                tot_scf_iters, delta_charge_sd, label=r"$\Delta$ charge/spin density"
            )

        # Add the convergence parameters
        if conv_params is not None and conv_params["charge_density"] != 0.0:
            ax.axhline(
                conv_params["charge_density"],
                ls="--",
                c="gray",
                label=r"$\rho$ convergence criterion",
            )

        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        ax.set_yscale("log")
        ax.set_xlabel("cumulative SCF iteration")
        ax.set_ylabel(r"Charge / $e a_0^{-3}$")
        ax.legend()
        if title is not None:
            ax.set_title(rf"{title} $\Delta$ charge Convergence")

        return ax

    @staticmethod
    def _plot_energy_convergence(
        ax: axes.Axes,
        tot_scf_iters: npt.NDArray[np.int64] | list[int],
        delta_sum_eigenvalues: npt.NDArray[np.float64] | list[float],
        delta_total_energies: npt.NDArray[np.float64] | list[float],
        conv_params: dict | None = None,
        absolute: bool = True,
        title: str | None = None,
    ) -> axes.Axes:
        """
        Create a subplot for the energy convergence of an FHI-aims calculation.

        Parameters
        ----------
        ax : axes.Axes
            matplotlib subplot index
        tot_scf_iters : npt.NDArray[np.int64] | list[int]
            cumulative SCF iterations
        delta_sum_eigenvalues : npt.NDArray[np.float64] | list[float]
            change of sum of eigenvalues
        delta_total_energies : npt.NDArray[np.float64] | list[float]
            change of total energies
        conv_params : dict | None, default=None
            convergence parameters which determine if the SCF cycle has converged
        absolute : bool, default=True
            whether to plot the absolute value of the energies
        title : str | None, default=None
            system name to include in title

        Returns
        -------
        axes.Axes
            matplotlib subplot object
        """
        if absolute:
            delta_sum_eigenvalues = np.abs(delta_sum_eigenvalues)
            delta_total_energies = np.abs(delta_total_energies)

        ax.plot(
            tot_scf_iters,
            delta_sum_eigenvalues,
            label=r"$\Delta \; \Sigma$ eigenvalues",
            c="C1",
        )
        ax.plot(
            tot_scf_iters,
            delta_total_energies,
            label=r"$\Delta$ total energies",
            c="C0",
        )

        # Add the convergence parameters
        if conv_params is not None:
            ax.axhline(
                conv_params["sum_eigenvalues"],
                ls="-.",
                c="darkgray",
                label=r"$\Delta \; \Sigma$ eigenvalues convergence criterion",
            )
            ax.axhline(
                conv_params["total_energy"],
                ls="--",
                c="gray",
                label=r"$\Delta$ total energies convergence criterion",
            )

        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        ax.set_yscale("log")
        ax.set_xlabel("cumulative SCF iteration")
        ax.set_ylabel(r"absolute energy / $| \mathrm{eV} |$")
        ax.legend()
        if title is not None:
            ax.set_title(rf"{title} Energies and Eigenvalues Convergence")

        return ax

    @staticmethod
    def _plot_forces_convergence(
        ax: axes.Axes,
        forces_on_atoms: npt.NDArray[np.float64] | list[float],
        conv_params: dict | None = None,
        title: str | None = None,
    ) -> None:
        """
        Create a subplot for the forces convergence of an FHI-aims calculation.

        Parameters
        ----------
        ax : axes.Axes
            matplotlib subplot index
        forces_on_atoms : npt.NDArray[np.float64] | list[float]
            all forces acting on each atom
        conv_params : dict | None, default=None
            convergence parameters which determine if the SCF cycle has converged
        title : str | None, default=None
            system name to include in title
        """
        # see NOTE in dfttoolkit.output.AimsOutput.get_i_scf_conv_acc()
        # ax.plot(delta_forces, label=r"$\Delta$ forces")  # noqa: ERA001
        ax.plot(forces_on_atoms, label="forces on atoms")

        # Add the convergence parameters
        if conv_params is not None and conv_params["total_force"] is not None:
            ax.axhline(
                conv_params["total_force"],
                ls="--",
                c="gray",
                label="forces convergence criterion",
            )

        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        ax.set_xlabel("geometry relaxation step")
        ax.set_ylabel(r"force / $\mathrm{eV} \mathrm{\AA}^{-1}$")
        ax.legend()

        if title is not None:
            ax.set_title(rf"{title} Forces Convergence")

    @staticmethod
    def _plot_ks_states_convergence(
        ax: axes.Axes,
        ks_eigenvals: dict | tuple[npt.NDArray, npt.NDArray],
        title: str | None = None,
    ) -> None:
        """
        Create a subplot for the energy changes of the KS eigenstates.

        Parameters
        ----------
        ax : axes.Axes
            matplotlib subplot index
        ks_eigenvals : dict | tuple[npt.NDArray, npt.NDArray]
            state, occupation, and eigenvalue of each KS state at each SCF
            iteration
        title : str | None, default=None
            system name to include in title

        Returns
        -------
        axes.Axes
            matplotlib subplot object
        """
        if isinstance(ks_eigenvals, dict):
            # Don't include last eigenvalue as it only prints after final SCF iteration
            # Add 1 to total SCF iterations to match the length of the eigenvalues and
            # we want to include the first pre SCF iteration
            for ev in ks_eigenvals["eigenvalue_eV"].T:
                ax.plot(np.arange(len(ks_eigenvals["eigenvalue_eV"])), ev)

        elif isinstance(ks_eigenvals, tuple):
            su_ks_eigenvals = ks_eigenvals[0]
            sd_ks_eigenvals = ks_eigenvals[1]

            for ev in su_ks_eigenvals["eigenvalue_eV"].T:
                ax.plot(np.arange(len(su_ks_eigenvals["eigenvalue_eV"])), ev, c="C0")

            for ev in sd_ks_eigenvals["eigenvalue_eV"].T:
                ax.plot(np.arange(len(su_ks_eigenvals["eigenvalue_eV"])), ev, c="C1")

        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        ax.set_yscale("symlog")
        ax.set_xlabel("cumulative SCF iteration")
        ax.set_ylabel("energy / eV")

        if title is not None:
            ax.set_title(f"{title} KS State Convergence")

    def convergence(
        self,
        conv_params: dict | None = None,
        scf_conv_acc_params: dict | None = None,
        title: str | None = None,
        forces: bool = False,
        ks_eigenvalues: bool = False,
        fig_size: tuple[int, int] = (24, 6),
    ) -> figure.Figure:
        """
        Plot the SCF convergence accuracy parameters.

        Parameters
        ----------
        conv_params: dict | None, default=None
            convergence parameters which determine if the SCF cycle has converged
        scf_conv_acc_params : dict | None, default=None
            the scf convergence accuracy parameters
        title : str | None, default=None
            system name to use in title of the plot
        forces : bool, default=False
            whether to plot the change of forces and forces on atoms
        ks_eigenvalues : bool, default=False
            whether to plot the kohn-sham eigenvalues
        fig_size : tuple[int, int], default=(24, 6)
            the total size of the figure

        Returns
        -------
        figure.Figure
            matplotlib figure object
        """
        # Get the SCF convergence accuracy parameters if not provided
        if scf_conv_acc_params is None:
            if not hasattr(self, "scf_conv_acc_params"):
                self.scf_conv_acc_params = self.get_i_scf_conv_acc()

        # Override the default scf_conv_acc_params if given in function
        else:
            self.scf_conv_acc_params = scf_conv_acc_params

        scf_iters = self.scf_conv_acc_params["scf_iter"]
        tot_scf_iters = np.arange(1, len(scf_iters) + 1)
        delta_charge = self.scf_conv_acc_params["change_of_charge"]
        delta_charge_sd = self.scf_conv_acc_params["change_of_charge_spin_density"]
        delta_sum_eigenvalues = self.scf_conv_acc_params["change_of_sum_eigenvalues"]
        delta_total_energies = self.scf_conv_acc_params["change_of_total_energy"]

        # Change the number of subplots if forces and ks_eigenvalues are to be plotted
        subplots = [True, True, forces, ks_eigenvalues]
        i_subplot = 1

        # Setup the figure subplots
        fig, ax = plt.subplots(1, subplots.count(True), figsize=fig_size)

        # Plot the change of charge
        self._plot_charge_convergence(
            ax[0], tot_scf_iters, delta_charge, delta_charge_sd, conv_params, title
        )

        # Plot the change of total energies and sum of eigenvalues
        self._plot_energy_convergence(
            ax[1],
            tot_scf_iters,
            delta_sum_eigenvalues,
            delta_total_energies,
            conv_params,
            True,
            title,
        )

        # Plot the forces
        if forces:
            # see NOTE in dfttoolkit.output.AimsOutput.get_i_scf_conv_acc()
            # delta_forces = self.scf_conv_acc_params["change_of_forces"]  # noqa: E501, ERA001
            # delta_forces = np.delete(delta_forces, np.argwhere(delta_forces == 0.0))  # noqa: E501, ERA001
            forces_on_atoms = self.scf_conv_acc_params["forces_on_atoms"]
            forces_on_atoms = np.delete(
                forces_on_atoms, np.argwhere(forces_on_atoms == 0.0)
            )
            i_subplot += 1
            self._plot_forces_convergence(
                ax[i_subplot], forces_on_atoms, conv_params, title
            )

        # Plot the KS state energies
        if ks_eigenvalues:
            i_subplot += 1
            ks_eigenvals = self.get_all_ks_eigenvalues()

            self._plot_ks_states_convergence(ax[i_subplot], ks_eigenvals, title)

        return fig

1	import matplotlib.pyplot as plt	×
2	import numpy as np	×
3	import numpy.typing as npt	×
4	from matplotlib import axes, figure	×
5	from matplotlib.ticker import MaxNLocator	×
6
7	from dfttoolkit.output import AimsOutput	×
8
9
10	class VisualiseAims(AimsOutput):	×
11	"""
12	FHI-aims visualisation tools.
13
14	...
15
16	Attributes
17	----------
18	scf_conv_acc_params : dict
19	the SCF convergence accuracy parameters
20	"""
21
22	@staticmethod	×
23	def _plot_charge_convergence(	×
24	ax: axes.Axes,
25	tot_scf_iters: npt.NDArray[np.int64] \| list[int],
26	delta_charge: npt.NDArray[np.float64] \| list[float],
27	delta_charge_sd: npt.NDArray[np.float64] \| list[float] \| None = None,
28	conv_params: dict \| None = None,
29	title: str \| None = None,
30	) -> axes.Axes:
31	"""
32	Create a subplot for the charge convergence of an FHI-aims calculation.
33
34	Parameters
35	----------
36	ax : axes.Axes
37	matplotlib subplot index
38	tot_scf_iters : npt.NDArray[np.int64] \| list[int]
39	cumulative SCF iterations
40	delta_charge : npt.NDArray[np.float64] \| list[float]
41	change of spin-up or total spin (if the calculation was spin none)
42	eigenvalues
43	delta_charge_sd : npt.NDArray[np.float64] \| list[float] \| None, default=None
44	change of spin-down eigenvalues
45	conv_params : dict \| None, default=None
46	convergence parameters which determine if the SCF cycle has converged
47	title : str \| None, default=None
48	system name to include in title
49
50	Returns
51	-------
52	axes.Axes
53	matplotlib subplot object
54	"""
55	ax.plot(tot_scf_iters, delta_charge, label=r"$\Delta$ charge")	×
56
57	# Only plot delta_charge_sd if the calculation is spin polarised
58	if delta_charge_sd is not None:	×
59	ax.plot(	×
60	tot_scf_iters, delta_charge_sd, label=r"$\Delta$ charge/spin density"
61	)
62
63	# Add the convergence parameters
64	if conv_params is not None and conv_params["charge_density"] != 0.0:	×
65	ax.axhline(	×
66	conv_params["charge_density"],
67	ls="--",
68	c="gray",
69	label=r"$\rho$ convergence criterion",
70	)
71
72	ax.xaxis.set_major_locator(MaxNLocator(integer=True))	×
73	ax.set_yscale("log")	×
74	ax.set_xlabel("cumulative SCF iteration")	×
75	ax.set_ylabel(r"Charge / $e a_0^{-3}$")	×
76	ax.legend()	×
77	if title is not None:	×
78	ax.set_title(rf"{title} $\Delta$ charge Convergence")	×
79
80	return ax	×
81
82	@staticmethod	×
83	def _plot_energy_convergence(	×
84	ax: axes.Axes,
85	tot_scf_iters: npt.NDArray[np.int64] \| list[int],
86	delta_sum_eigenvalues: npt.NDArray[np.float64] \| list[float],
87	delta_total_energies: npt.NDArray[np.float64] \| list[float],
88	conv_params: dict \| None = None,
89	absolute: bool = True,
90	title: str \| None = None,
91	) -> axes.Axes:
92	"""
93	Create a subplot for the energy convergence of an FHI-aims calculation.
94
95	Parameters
96	----------
97	ax : axes.Axes
98	matplotlib subplot index
99	tot_scf_iters : npt.NDArray[np.int64] \| list[int]
100	cumulative SCF iterations
101	delta_sum_eigenvalues : npt.NDArray[np.float64] \| list[float]
102	change of sum of eigenvalues
103	delta_total_energies : npt.NDArray[np.float64] \| list[float]
104	change of total energies
105	conv_params : dict \| None, default=None
106	convergence parameters which determine if the SCF cycle has converged
107	absolute : bool, default=True
108	whether to plot the absolute value of the energies
109	title : str \| None, default=None
110	system name to include in title
111
112	Returns
113	-------
114	axes.Axes
115	matplotlib subplot object
116	"""
117	if absolute:	×
118	delta_sum_eigenvalues = np.abs(delta_sum_eigenvalues)	×
119	delta_total_energies = np.abs(delta_total_energies)	×
120
121	ax.plot(	×
122	tot_scf_iters,
123	delta_sum_eigenvalues,
124	label=r"$\Delta \; \Sigma$ eigenvalues",
125	c="C1",
126	)
127	ax.plot(	×
128	tot_scf_iters,
129	delta_total_energies,
130	label=r"$\Delta$ total energies",
131	c="C0",
132	)
133
134	# Add the convergence parameters
135	if conv_params is not None:	×
136	ax.axhline(	×
137	conv_params["sum_eigenvalues"],
138	ls="-.",
139	c="darkgray",
140	label=r"$\Delta \; \Sigma$ eigenvalues convergence criterion",
141	)
142	ax.axhline(	×
143	conv_params["total_energy"],
144	ls="--",
145	c="gray",
146	label=r"$\Delta$ total energies convergence criterion",
147	)
148
149	ax.xaxis.set_major_locator(MaxNLocator(integer=True))	×
150	ax.set_yscale("log")	×
151	ax.set_xlabel("cumulative SCF iteration")	×
152	ax.set_ylabel(r"absolute energy / $\| \mathrm{eV} \|$")	×
153	ax.legend()	×
154	if title is not None:	×
155	ax.set_title(rf"{title} Energies and Eigenvalues Convergence")	×
156
157	return ax	×
158
159	@staticmethod	×
160	def _plot_forces_convergence(	×
161	ax: axes.Axes,
162	forces_on_atoms: npt.NDArray[np.float64] \| list[float],
163	conv_params: dict \| None = None,
164	title: str \| None = None,
165	) -> None:
166	"""
167	Create a subplot for the forces convergence of an FHI-aims calculation.
168
169	Parameters
170	----------
171	ax : axes.Axes
172	matplotlib subplot index
173	forces_on_atoms : npt.NDArray[np.float64] \| list[float]
174	all forces acting on each atom
175	conv_params : dict \| None, default=None
176	convergence parameters which determine if the SCF cycle has converged
177	title : str \| None, default=None
178	system name to include in title
179	"""
180	# see NOTE in dfttoolkit.output.AimsOutput.get_i_scf_conv_acc()
181	# ax.plot(delta_forces, label=r"$\Delta$ forces") # noqa: ERA001
182	ax.plot(forces_on_atoms, label="forces on atoms")	×
183
184	# Add the convergence parameters
185	if conv_params is not None and conv_params["total_force"] is not None:	×
186	ax.axhline(	×
187	conv_params["total_force"],
188	ls="--",
189	c="gray",
190	label="forces convergence criterion",
191	)
192
193	ax.xaxis.set_major_locator(MaxNLocator(integer=True))	×
194	ax.set_xlabel("geometry relaxation step")	×
195	ax.set_ylabel(r"force / $\mathrm{eV} \mathrm{\AA}^{-1}$")	×
196	ax.legend()	×
197
198	if title is not None:	×
199	ax.set_title(rf"{title} Forces Convergence")	×
200
201	@staticmethod	×
202	def _plot_ks_states_convergence(	×
203	ax: axes.Axes,
204	ks_eigenvals: dict \| tuple[npt.NDArray, npt.NDArray],
205	title: str \| None = None,
206	) -> None:
207	"""
208	Create a subplot for the energy changes of the KS eigenstates.
209
210	Parameters
211	----------
212	ax : axes.Axes
213	matplotlib subplot index
214	ks_eigenvals : dict \| tuple[npt.NDArray, npt.NDArray]
215	state, occupation, and eigenvalue of each KS state at each SCF
216	iteration
217	title : str \| None, default=None
218	system name to include in title
219
220	Returns
221	-------
222	axes.Axes
223	matplotlib subplot object
224	"""
225	if isinstance(ks_eigenvals, dict):	×
226	# Don't include last eigenvalue as it only prints after final SCF iteration
227	# Add 1 to total SCF iterations to match the length of the eigenvalues and
228	# we want to include the first pre SCF iteration
229	for ev in ks_eigenvals["eigenvalue_eV"].T:	×
230	ax.plot(np.arange(len(ks_eigenvals["eigenvalue_eV"])), ev)	×
231
232	elif isinstance(ks_eigenvals, tuple):	×
233	su_ks_eigenvals = ks_eigenvals[0]	×
234	sd_ks_eigenvals = ks_eigenvals[1]	×
235
236	for ev in su_ks_eigenvals["eigenvalue_eV"].T:	×
237	ax.plot(np.arange(len(su_ks_eigenvals["eigenvalue_eV"])), ev, c="C0")	×
238
239	for ev in sd_ks_eigenvals["eigenvalue_eV"].T:	×
240	ax.plot(np.arange(len(su_ks_eigenvals["eigenvalue_eV"])), ev, c="C1")	×
241
242	ax.xaxis.set_major_locator(MaxNLocator(integer=True))	×
243	ax.set_yscale("symlog")	×
244	ax.set_xlabel("cumulative SCF iteration")	×
245	ax.set_ylabel("energy / eV")	×
246
247	if title is not None:	×
248	ax.set_title(f"{title} KS State Convergence")	×
249
250	def convergence(	×
251	self,
252	conv_params: dict \| None = None,
253	scf_conv_acc_params: dict \| None = None,
254	title: str \| None = None,
255	forces: bool = False,
256	ks_eigenvalues: bool = False,
257	fig_size: tuple[int, int] = (24, 6),
258	) -> figure.Figure:
259	"""
260	Plot the SCF convergence accuracy parameters.
261
262	Parameters
263	----------
264	conv_params: dict \| None, default=None
265	convergence parameters which determine if the SCF cycle has converged
266	scf_conv_acc_params : dict \| None, default=None
267	the scf convergence accuracy parameters
268	title : str \| None, default=None
269	system name to use in title of the plot
270	forces : bool, default=False
271	whether to plot the change of forces and forces on atoms
272	ks_eigenvalues : bool, default=False
273	whether to plot the kohn-sham eigenvalues
274	fig_size : tuple[int, int], default=(24, 6)
275	the total size of the figure
276
277	Returns
278	-------
279	figure.Figure
280	matplotlib figure object
281	"""
282	# Get the SCF convergence accuracy parameters if not provided
283	if scf_conv_acc_params is None:	×
284	if not hasattr(self, "scf_conv_acc_params"):	×
285	self.scf_conv_acc_params = self.get_i_scf_conv_acc()	×
286
287	# Override the default scf_conv_acc_params if given in function
288	else:
289	self.scf_conv_acc_params = scf_conv_acc_params	×
290
291	scf_iters = self.scf_conv_acc_params["scf_iter"]	×
292	tot_scf_iters = np.arange(1, len(scf_iters) + 1)	×
293	delta_charge = self.scf_conv_acc_params["change_of_charge"]	×
294	delta_charge_sd = self.scf_conv_acc_params["change_of_charge_spin_density"]	×
295	delta_sum_eigenvalues = self.scf_conv_acc_params["change_of_sum_eigenvalues"]	×
296	delta_total_energies = self.scf_conv_acc_params["change_of_total_energy"]	×
297
298	# Change the number of subplots if forces and ks_eigenvalues are to be plotted
299	subplots = [True, True, forces, ks_eigenvalues]	×
300	i_subplot = 1	×
301
302	# Setup the figure subplots
303	fig, ax = plt.subplots(1, subplots.count(True), figsize=fig_size)	×
304
305	# Plot the change of charge
306	self._plot_charge_convergence(	×
307	ax[0], tot_scf_iters, delta_charge, delta_charge_sd, conv_params, title
308	)
309
310	# Plot the change of total energies and sum of eigenvalues
311	self._plot_energy_convergence(	×
312	ax[1],
313	tot_scf_iters,
314	delta_sum_eigenvalues,
315	delta_total_energies,
316	conv_params,
317	True,
318	title,
319	)
320
321	# Plot the forces
322	if forces:	×
323	# see NOTE in dfttoolkit.output.AimsOutput.get_i_scf_conv_acc()
324	# delta_forces = self.scf_conv_acc_params["change_of_forces"] # noqa: E501, ERA001
325	# delta_forces = np.delete(delta_forces, np.argwhere(delta_forces == 0.0)) # noqa: E501, ERA001
326	forces_on_atoms = self.scf_conv_acc_params["forces_on_atoms"]	×
327	forces_on_atoms = np.delete(	×
328	forces_on_atoms, np.argwhere(forces_on_atoms == 0.0)
329	)
330	i_subplot += 1	×
331	self._plot_forces_convergence(	×
332	ax[i_subplot], forces_on_atoms, conv_params, title
333	)
334
335	# Plot the KS state energies
336	if ks_eigenvalues:	×
337	i_subplot += 1	×
338	ks_eigenvals = self.get_all_ks_eigenvalues()	×
339
340	self._plot_ks_states_convergence(ax[i_subplot], ks_eigenvals, title)	×
341
342	return fig	×

maurergroup / dfttoolkit / 15077298886

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