14995306316

Committed 13 May 2025 11:21AM UTC coverage: 83.844%. Remained the same

Build # 14995306316

Build Type

Pull #194

github

Committed by

web-flow

Commit Message

Merge d8ac4723f into dd0c5094e

Pull Request Pull Request #194: apply black to code base

Run Details

955 of 1334 branches covered (71.59%)

Branch coverage included in aggregate %.

240 of 268 new or added lines in 50 files covered. (89.55%)

5 existing lines in 4 files now uncovered.

4515 of 5190 relevant lines covered (86.99%)

0.87 hits per line

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

55.93

/qmctorch/utils/plot_data.py

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
from types import SimpleNamespace
from typing import Optional, Union, Tuple
from .stat_utils import (
    blocking,
    correlation_coefficient,
    fit_correlation_coefficient,
    integrated_autocorrelation_time,
)


def plot_energy(
    local_energy: np.ndarray,
    e0: Optional[float] = None,
    show_variance: bool = False,
    clip: bool = False,
    q: float = 0.15,
) -> None:
    """Plot the evolution of the energy.

    Args:
        local_energy (np.ndarray): Local energies along the trajectory.
        e0 (float, optional): Target value for the energy. Defaults to None.
        show_variance (bool, optional): Show the variance if True. Defaults to False.
        clip (bool, optional): Clip the values to remove outliers. Defaults to False.
        q (float, optional): Quantile used for the interquartile range. Defaults to 0.15.
    """

    def clip_values(values: np.ndarray, std_factor: int = 5) -> np.ndarray:
        if clip:
            values = values.flatten()
            mean = np.median(values)
            std = values.std()
            up = values < mean + std_factor * std
            down = values > mean - std_factor * std
            return values[up * down]
        return values

    fig = plt.figure()
    ax = fig.add_subplot(111)

    n = len(local_energy)
    epoch = np.arange(n)

    # get the variance

    energy = np.array([np.mean(clip_values(e)) for e in local_energy])
    variance = np.array([np.var(clip_values(e)) for e in local_energy])
    q75 = np.array([np.quantile(clip_values(e), 0.5 + q) for e in local_energy])
    q25 = np.array([np.quantile(clip_values(e), 0.5 - q) for e in local_energy])

    # plot
    ax.fill_between(epoch, q25, q75, alpha=0.5, color="#4298f4")
    ax.plot(epoch, energy, color="#144477")
    if e0 is not None:
        ax.axhline(e0, color="black", linestyle="--")

    ax.grid()
    ax.set_xlabel("Number of epoch")
    ax.set_ylabel("Energy", color="black")

    if show_variance:
        ax2 = ax.twinx()
        ax2.plot(epoch, variance, color="blue")
        ax2.set_ylabel("variance", color="blue")
        ax2.tick_params(axis="y", labelcolor="blue")
        fig.tight_layout()

    plt.show()


def plot_data(observable: SimpleNamespace, obsname: str) -> None:
    """Plot the evolution of a given data

    Args:
        observable (SimpleNamespace): namespace of observable
        obsname (str): name (key) of the desired observable
    """

    fig = plt.figure()
    ax = fig.add_subplot(111)

    data = np.array(observable.__dict__[obsname]).squeeze()
    epoch = np.arange(len(data))
    ax.plot(epoch, data, color="#144477")

    plt.show()


def plot_walkers_traj(
    eloc: np.ndarray, walkers: Union[int, str, None] = "mean"
) -> None:
    """Plot the trajectory of all the individual walkers

    Args:
        eloc (np.ndarray): Local energy array (Nstep, Nwalkers)
        walkers (int, str, optional): all, mean or index of a given walker Defaults to 'mean'

    Returns:
        None
    """
    nstep, nwalkers = eloc.shape
    celoc = np.cumsum(eloc, axis=0).T
    celoc /= np.arange(1, nstep + 1)

    if walkers is not None:
        if walkers == "all":
            plt.plot(eloc, "o", alpha=max(1 / nwalkers, 1e-2), c="grey")
            cmap = cm.hot(np.linspace(0, 1, nwalkers))
            for i in range(nwalkers):
                plt.plot(celoc.T[:, i], color=cmap[i])

        elif walkers == "mean":
            plt.plot(eloc, "o", alpha=1 / nwalkers, c="grey")
            emean = np.mean(celoc.T, axis=1)
            emin = emean.min()
            emax = emean.max()
            delta = emax - emin
            plt.plot(emean, linewidth=5)
            plt.ylim(emin - 0.25 * delta, emax + 0.25 * delta)
        else:
            raise ValueError("walkers argument must be all or mean")

        plt.grid()
        plt.xlabel("Monte Carlo Steps")
        plt.ylabel("Energy (Hartree)")

    plt.show()


def plot_correlation_coefficient(
    eloc: np.ndarray, size_max: int = 100
) -> Tuple[np.ndarray, float]:
    """
    Plot the correlation coefficient of the local energy
    and fit the curve to an exp to extract the correlation time.

    Parameters
    ----------
    eloc : np.ndarray
        values of the local energy (Nstep, Nwalk)
    size_max : int, optional
        maximu number of MC step to consider. Defaults to 100.

    Returns
    -------
    rho : np.ndarray
        correlation coefficients (size_max, Nwalkers)
    tau_fit : float
        correlation time
    """
    rho = correlation_coefficient(eloc)
    tau_fit, fitted = fit_correlation_coefficient(rho.mean(1)[:size_max])

    plt.plot(rho, alpha=0.25)
    plt.plot(rho.mean(1), linewidth=3, c="black")
    plt.plot(fitted, "--", c="grey")
    plt.xlim([0, size_max])
    plt.ylim([-0.25, 1.5])
    plt.xlabel("MC steps")
    plt.ylabel("Correlation coefficient")
    plt.text(
        0.5 * size_max, 1.05, "tau=%1.3f" % tau_fit, {"color": "black", "fontsize": 15}
    )
    plt.grid()
    plt.show()

    return rho, tau_fit


def plot_integrated_autocorrelation_time(
    eloc: np.ndarray, rho: np.ndarray = None, size_max: int = 100, C: int = 5
) -> int:
    """Compute and plot the integrated autocorrelation time.

    Args:
        eloc (np.ndarray): Local energy values (Nstep, Nwalkers).
        rho (np.ndarray, optional): Correlation coefficient. Defaults to None.
        size_max (int, optional): Maximum number of MC steps to consider. Defaults to 100.
        C (int, optional): A constant used for thresholding. Defaults to 5.

    Returns:
        int: Index where the mean integrated autocorrelation time meets the condition.
    """
    if rho is None:
        rho = correlation_coefficient(eloc)

    tau = integrated_autocorrelation_time(rho, size_max)

    tc, idx_tc = [], []
    idx = np.arange(1, size_max)
    for iw in range(eloc.shape[1]):
        t = tau[:, iw]
        if len(t[C * t <= idx]) > 0:
            tval = t[C * t <= idx][0]
            ii = np.where(t == tval)[0][0]

            tc.append(tval)
            idx_tc.append(ii)

    plt.plot(tau, alpha=0.25)
    tm = tau.mean(1)
    plt.plot(tm, c="black")
    plt.plot(idx / C, "--", c="grey")

    plt.plot(idx_tc, tc, "o", alpha=0.25)
    tt = tm[tm * C <= idx][0]
    ii = np.where(tm == tt)[0][0]
    plt.plot(ii, tt, "o")

    plt.grid()
    plt.xlabel("MC step")
    plt.ylabel("IAC")
    plt.show()

    return ii


def plot_blocking_energy(
    eloc: np.ndarray, block_size: int, walkers: str = "mean"
) -> np.ndarray:
    """Plot the blocked energy values

    Args:
        eloc (np.ndarray): values of the local energies
        block_size (int): size of the block
        walkers (str, optional): which walkers to plot (mean, all, index or list). Defaults to 'mean'.

    Returns:
        np.ndarray: blocked energy values

    Raises:
        ValueError: [description]
    """
    eb = blocking(eloc, block_size, expand=True)
    if walkers == "all":
        plt.plot(eloc)
        plt.plot(eb)

    elif walkers == "mean":
        plt.plot(eloc.mean(1))
        plt.plot(eb.mean(1))

    elif walkers.__class__.__name__ in ["int", "list"]:
        plt.plot(eloc[:, walkers])
        plt.plot(eb[:, walkers])

    else:
        raise ValueError("walkers ", walkers, " not recognized")

    plt.grid()
    plt.xlabel("MC steps")
    plt.ylabel("Energy")
    plt.show()

    return blocking(eloc, block_size, expand=False)


def plot_correlation_time(eloc: np.ndarray) -> None:
    """Plot the blocking thingy

    Args:
        eloc (np.ndarray): values of the local energy
    """

    nstep, _ = eloc.shape
    max_block_size = nstep // 2

    var = np.std(eloc, axis=0)

    evar = []
    for size in range(1, max_block_size):
        eb = blocking(eloc, size)
        evar.append(np.std(eb, axis=0) * size / var)

    plt.plot(np.array(evar))
    plt.xlabel("Blocking size")
    plt.ylabel("Correlation steps")
    plt.show()


def plot_block(eloc: np.ndarray) -> None:
    """Plot the standard error of the blocked energies.

    Args:
        eloc (np.ndarray): Values of the local energy.

    Returns:
        None
    """

    nstep, _ = eloc.shape
    max_block_size = nstep // 2

    evar = []
    for size in range(1, max_block_size):
        eb = blocking(eloc, size)
        nblock = eb.shape[0]
        evar.append(np.sqrt(np.var(eb, axis=0) / (nblock - 1)))

    plt.plot(np.array(evar))
    plt.show()

1	import matplotlib.pyplot as plt	1✔
2	import numpy as np	1✔
3	from matplotlib import cm	1✔
4	from types import SimpleNamespace	1✔
5	from typing import Optional, Union, Tuple	1✔
6	from .stat_utils import (	1✔
7	blocking,
8	correlation_coefficient,
9	fit_correlation_coefficient,
10	integrated_autocorrelation_time,
11	)
12
13
14	def plot_energy(	1✔
15	local_energy: np.ndarray,
16	e0: Optional[float] = None,
17	show_variance: bool = False,
18	clip: bool = False,
19	q: float = 0.15,
20	) -> None:
21	"""Plot the evolution of the energy.
22
23	Args:
24	local_energy (np.ndarray): Local energies along the trajectory.
25	e0 (float, optional): Target value for the energy. Defaults to None.
26	show_variance (bool, optional): Show the variance if True. Defaults to False.
27	clip (bool, optional): Clip the values to remove outliers. Defaults to False.
28	q (float, optional): Quantile used for the interquartile range. Defaults to 0.15.
29	"""
30
31	def clip_values(values: np.ndarray, std_factor: int = 5) -> np.ndarray:	×
32	if clip:	×
33	values = values.flatten()	×
34	mean = np.median(values)	×
35	std = values.std()	×
36	up = values < mean + std_factor * std	×
37	down = values > mean - std_factor * std	×
38	return values[up * down]	×
39	return values	×
40
41	fig = plt.figure()	×
42	ax = fig.add_subplot(111)	×
43
44	n = len(local_energy)	×
45	epoch = np.arange(n)	×
46
47	# get the variance
48
49	energy = np.array([np.mean(clip_values(e)) for e in local_energy])	×
50	variance = np.array([np.var(clip_values(e)) for e in local_energy])	×
51	q75 = np.array([np.quantile(clip_values(e), 0.5 + q) for e in local_energy])	×
52	q25 = np.array([np.quantile(clip_values(e), 0.5 - q) for e in local_energy])	×
53
54	# plot
NEW 55	ax.fill_between(epoch, q25, q75, alpha=0.5, color="#4298f4")	×
UNCOV 56	ax.plot(epoch, energy, color="#144477")	×
57	if e0 is not None:	×
58	ax.axhline(e0, color="black", linestyle="--")	×
59
60	ax.grid()	×
61	ax.set_xlabel("Number of epoch")	×
62	ax.set_ylabel("Energy", color="black")	×
63
64	if show_variance:	×
65	ax2 = ax.twinx()	×
66	ax2.plot(epoch, variance, color="blue")	×
67	ax2.set_ylabel("variance", color="blue")	×
68	ax2.tick_params(axis="y", labelcolor="blue")	×
69	fig.tight_layout()	×
70
71	plt.show()	×
72
73
74	def plot_data(observable: SimpleNamespace, obsname: str) -> None:	1✔
75	"""Plot the evolution of a given data
76
77	Args:
78	observable (SimpleNamespace): namespace of observable
79	obsname (str): name (key) of the desired observable
80	"""
81
82	fig = plt.figure()	×
83	ax = fig.add_subplot(111)	×
84
85	data = np.array(observable.__dict__[obsname]).squeeze()	×
86	epoch = np.arange(len(data))	×
87	ax.plot(epoch, data, color="#144477")	×
88
89	plt.show()	×
90
91
92	def plot_walkers_traj(	1✔
93	eloc: np.ndarray, walkers: Union[int, str, None] = "mean"
94	) -> None:
95	"""Plot the trajectory of all the individual walkers
96
97	Args:
98	eloc (np.ndarray): Local energy array (Nstep, Nwalkers)
99	walkers (int, str, optional): all, mean or index of a given walker Defaults to 'mean'
100
101	Returns:
102	None
103	"""
104	nstep, nwalkers = eloc.shape	1✔
105	celoc = np.cumsum(eloc, axis=0).T	1✔
106	celoc /= np.arange(1, nstep + 1)	1✔
107
108	if walkers is not None:	1!
109	if walkers == "all":	1!
110	plt.plot(eloc, "o", alpha=max(1 / nwalkers, 1e-2), c="grey")	×
111	cmap = cm.hot(np.linspace(0, 1, nwalkers))	×
112	for i in range(nwalkers):	×
113	plt.plot(celoc.T[:, i], color=cmap[i])	×
114
115	elif walkers == "mean":	1!
116	plt.plot(eloc, "o", alpha=1 / nwalkers, c="grey")	1✔
117	emean = np.mean(celoc.T, axis=1)	1✔
118	emin = emean.min()	1✔
119	emax = emean.max()	1✔
120	delta = emax - emin	1✔
121	plt.plot(emean, linewidth=5)	1✔
122	plt.ylim(emin - 0.25 * delta, emax + 0.25 * delta)	1✔
123	else:
124	raise ValueError("walkers argument must be all or mean")	×
125
126	plt.grid()	1✔
127	plt.xlabel("Monte Carlo Steps")	1✔
128	plt.ylabel("Energy (Hartree)")	1✔
129
130	plt.show()	1✔
131
132
133	def plot_correlation_coefficient(	1✔
134	eloc: np.ndarray, size_max: int = 100
135	) -> Tuple[np.ndarray, float]:
136	"""
137	Plot the correlation coefficient of the local energy
138	and fit the curve to an exp to extract the correlation time.
139
140	Parameters
141	----------
142	eloc : np.ndarray
143	values of the local energy (Nstep, Nwalk)
144	size_max : int, optional
145	maximu number of MC step to consider. Defaults to 100.
146
147	Returns
148	-------
149	rho : np.ndarray
150	correlation coefficients (size_max, Nwalkers)
151	tau_fit : float
152	correlation time
153	"""
154	rho = correlation_coefficient(eloc)	1✔
155	tau_fit, fitted = fit_correlation_coefficient(rho.mean(1)[:size_max])	1✔
156
157	plt.plot(rho, alpha=0.25)	1✔
158	plt.plot(rho.mean(1), linewidth=3, c="black")	1✔
159	plt.plot(fitted, "--", c="grey")	1✔
160	plt.xlim([0, size_max])	1✔
161	plt.ylim([-0.25, 1.5])	1✔
162	plt.xlabel("MC steps")	1✔
163	plt.ylabel("Correlation coefficient")	1✔
164	plt.text(	1✔
165	0.5 * size_max, 1.05, "tau=%1.3f" % tau_fit, {"color": "black", "fontsize": 15}
166	)
167	plt.grid()	1✔
168	plt.show()	1✔
169
170	return rho, tau_fit	1✔
171
172
173	def plot_integrated_autocorrelation_time(	1✔
174	eloc: np.ndarray, rho: np.ndarray = None, size_max: int = 100, C: int = 5
175	) -> int:
176	"""Compute and plot the integrated autocorrelation time.
177
178	Args:
179	eloc (np.ndarray): Local energy values (Nstep, Nwalkers).
180	rho (np.ndarray, optional): Correlation coefficient. Defaults to None.
181	size_max (int, optional): Maximum number of MC steps to consider. Defaults to 100.
182	C (int, optional): A constant used for thresholding. Defaults to 5.
183
184	Returns:
185	int: Index where the mean integrated autocorrelation time meets the condition.
186	"""
187	if rho is None:	1!
188	rho = correlation_coefficient(eloc)	1✔
189
190	tau = integrated_autocorrelation_time(rho, size_max)	1✔
191
192	tc, idx_tc = [], []	1✔
193	idx = np.arange(1, size_max)	1✔
194	for iw in range(eloc.shape[1]):	1✔
195	t = tau[:, iw]	1✔
196	if len(t[C * t <= idx]) > 0:	1!
197	tval = t[C * t <= idx][0]	1✔
198	ii = np.where(t == tval)[0][0]	1✔
199
200	tc.append(tval)	1✔
201	idx_tc.append(ii)	1✔
202
203	plt.plot(tau, alpha=0.25)	1✔
204	tm = tau.mean(1)	1✔
205	plt.plot(tm, c="black")	1✔
206	plt.plot(idx / C, "--", c="grey")	1✔
207
208	plt.plot(idx_tc, tc, "o", alpha=0.25)	1✔
209	tt = tm[tm * C <= idx][0]	1✔
210	ii = np.where(tm == tt)[0][0]	1✔
211	plt.plot(ii, tt, "o")	1✔
212
213	plt.grid()	1✔
214	plt.xlabel("MC step")	1✔
215	plt.ylabel("IAC")	1✔
216	plt.show()	1✔
217
218	return ii	1✔
219
220
221	def plot_blocking_energy(	1✔
222	eloc: np.ndarray, block_size: int, walkers: str = "mean"
223	) -> np.ndarray:
224	"""Plot the blocked energy values
225
226	Args:
227	eloc (np.ndarray): values of the local energies
228	block_size (int): size of the block
229	walkers (str, optional): which walkers to plot (mean, all, index or list). Defaults to 'mean'.
230
231	Returns:
232	np.ndarray: blocked energy values
233
234	Raises:
235	ValueError: [description]
236	"""
237	eb = blocking(eloc, block_size, expand=True)	1✔
238	if walkers == "all":	1!
239	plt.plot(eloc)	×
240	plt.plot(eb)	×
241
242	elif walkers == "mean":	1!
243	plt.plot(eloc.mean(1))	1✔
244	plt.plot(eb.mean(1))	1✔
245
246	elif walkers.__class__.__name__ in ["int", "list"]:	×
247	plt.plot(eloc[:, walkers])	×
248	plt.plot(eb[:, walkers])	×
249
250	else:
251	raise ValueError("walkers ", walkers, " not recognized")	×
252
253	plt.grid()	1✔
254	plt.xlabel("MC steps")	1✔
255	plt.ylabel("Energy")	1✔
256	plt.show()	1✔
257
258	return blocking(eloc, block_size, expand=False)	1✔
259
260
261	def plot_correlation_time(eloc: np.ndarray) -> None:	1✔
262	"""Plot the blocking thingy
263
264	Args:
265	eloc (np.ndarray): values of the local energy
266	"""
267
268	nstep, _ = eloc.shape	×
269	max_block_size = nstep // 2	×
270
271	var = np.std(eloc, axis=0)	×
272
273	evar = []	×
274	for size in range(1, max_block_size):	×
275	eb = blocking(eloc, size)	×
276	evar.append(np.std(eb, axis=0) * size / var)	×
277
278	plt.plot(np.array(evar))	×
279	plt.xlabel("Blocking size")	×
280	plt.ylabel("Correlation steps")	×
281	plt.show()	×
282
283
284	def plot_block(eloc: np.ndarray) -> None:	1✔
285	"""Plot the standard error of the blocked energies.
286
287	Args:
288	eloc (np.ndarray): Values of the local energy.
289
290	Returns:
291	None
292	"""
293
294	nstep, _ = eloc.shape	1✔
295	max_block_size = nstep // 2	1✔
296
297	evar = []	1✔
298	for size in range(1, max_block_size):	1✔
299	eb = blocking(eloc, size)	1✔
300	nblock = eb.shape[0]	1✔
301	evar.append(np.sqrt(np.var(eb, axis=0) / (nblock - 1)))	1✔
302
303	plt.plot(np.array(evar))	1✔
304	plt.show()	1✔

NLESC-JCER / QMCTorch / 14995306316

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