14088074301

Committed 26 Mar 2025 04:07PM UTC coverage: 79.879% (+0.2%) from 79.686%

Build # 14088074301

Build Type

Pull #180

github

Committed by

web-flow

Commit Message

Merge b3c2c0b53 into b8624a46e

Pull Request Pull Request #180: Typehints

Run Details

1056 of 1570 branches covered (67.26%)

Branch coverage included in aggregate %.

596 of 634 new or added lines in 65 files covered. (94.01%)

6 existing lines in 6 files now uncovered.

5022 of 6039 relevant lines covered (83.16%)

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

NLESC-JCER / QMCTorch / 14088074301

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