Coverage for calorine/tools/entropy.py: 100%
71 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-14 16:36 +0000
1from warnings import warn
2import numpy as np
4try:
5 import torch
6 torch_available = True
7except ImportError: # pragma: no cover
8 torch_available = False
9 # If pytorch is not available, we still need a torch object
10 # to satisfy type hints etc. Otherwise the file will not run.
12 class _DummyTensor:
13 """Fallback Tensor placeholder."""
14 pass
16 class _DummyDevice:
17 """Fallback device placeholder."""
18 def __init__(self, *args, **kwargs):
19 pass
21 class _DummyDType:
22 """Fallback dtype placeholder."""
23 pass
25 class _DummyTorchModule:
26 Tensor = _DummyTensor
27 device = _DummyDevice
28 dtype = _DummyDType
30 # Provide common dtype names you might reference
31 float32 = _DummyDType()
32 float64 = _DummyDType()
34 torch = _DummyTorchModule()
36try:
37 import warnings as _warnings
38 with _warnings.catch_warnings():
39 _warnings.filterwarnings('ignore', message='IProgress not found')
40 from tqdm.auto import tqdm
41except ImportError:
42 def tqdm(iterable, **kwargs):
43 """Fallback for when tqdm is not installed; accepts and ignores the
44 keyword arguments that tqdm takes."""
45 return iterable
48def get_entropy(
49 descriptors: np.ndarray | torch.Tensor,
50 width: float,
51 use_tqdm: bool = True,
52 device: str | torch.device = 'cuda',
53 block: int = 1024,
54 dtype: torch.dtype = torch.float64,
55 eps: float = 1e-12,
56 cumulative: bool = False,
57) -> tuple[float, np.ndarray] | tuple[np.ndarray, np.ndarray]:
58 r"""
59 Computes an estimate for the information entropy :math:`H(\mathbf{X})` of a
60 set of descriptors :math:`\mathbf{X}`. The estimate is described in
61 [Nat. Comm. **16**, 4014 (2025)](https://doi.org/10.1038/s41467-025-59232-0)
62 and given by
64 .. math::
66 \mathcal{H}(\{\mathbf{X}\}) = -\frac{1}{n} \sum_{i=1}^{n} p_i
68 where
70 .. math::
72 p_i
73 = \log \left[
74 \frac{1}{n} \sum_{j=1}^{n}
75 K_h(\mathbf{X}_i, \mathbf{X}_j)
76 \right]
78 with a Gaussian kernel
80 .. math::
82 K_h(\mathbf{X}_i, \mathbf{X}_j)
83 = \exp\!\left(
84 -\frac{\lVert \mathbf{X}_i - \mathbf{X}_j \rVert^2}{2h^2}.
85 \right)
87 The calculation is done via torch if the library has been installed,
88 and numpy otherwise. When using torch the calculation is run via CUDA.
89 The latter behavior can be controlled using the :attr:`device` argument.
91 Parameters
92 ----------
93 descriptors
94 The set of descriptors :math:`\mathbf{X}` for which to evaluate to
95 the entropy. Typically each row corresponds to one atom and the
96 columns correspond to the different descriptor components.
97 width
98 Width :math:`h` of the Gaussian kernel.
99 use_tqdm
100 Use `tqdm <https://tqdm.github.io/>`_ to show a progress bar.
101 Note that this requires tqdm to be installed.
102 block
103 In order to limit the memory needs, the kernel density estimate
104 matrix is handled in blocks. This parameter controls the size of
105 each block. Smaller numbers imply a smaller memory footprint.
106 eps
107 Smallest (absolute) permissible value.
108 device
109 Device to use for calculation. The documentation of
110 [`torch.device`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device)
111 provides more information.
112 Only used when pytorch is available.
113 dtype
114 Floating point precision used for the computation. The documentation of
115 [`torch.dtype`](https://docs.pytorch.org/docs/stable/tensor_attributes.html)
116 provides an overview.
117 Only used when pytorch is available.
118 cumulative
119 If True, returns the cumulative entropy of each row in the descriptor.
122 Returns
123 -------
124 A tuple comprising the total entropy :math:`H(\mathbf{X})` and the entropy contributions
125 :math:`p_i` from each row in the input descriptor matrix.
126 """
127 if torch_available:
128 res = _get_entropy_torch(descriptors,
129 width,
130 use_tqdm,
131 block,
132 eps,
133 cumulative,
134 device,
135 dtype)
136 return res
137 else: # pragma: no cover
138 warn('Using the numpy implementation.'
139 ' Install torch in order to use GPUs and achieve a considerable speed-up.')
140 res = _get_entropy_numpy(descriptors, width, use_tqdm, block, eps, cumulative)
141 return res
144def _get_entropy_numpy(
145 descriptors: np.ndarray,
146 width: float,
147 use_tqdm: bool,
148 block: int,
149 eps: float,
150 cumulative: bool,
151) -> tuple[float, np.ndarray] | tuple[np.ndarray, np.ndarray]:
152 """Compute the informational entropy using numpy.
153 See get_entropy for documentation.
154 """
155 X = np.asarray(descriptors, dtype=np.float64)
156 N, d = X.shape
157 s = np.sum(X * X, axis=1) # (N,)
158 inv_two_sigma2 = 1.0 / (2.0 * width * width)
160 if not cumulative:
161 N_mat = N
162 row_sums = np.zeros(N, dtype=np.float64)
163 else:
164 N_mat = np.repeat(np.arange(1, N+1, dtype=np.float64)[:, None], N, axis=1).T
165 row_sums = np.tril(N_mat, -1)
167 for i0 in tqdm(range(0, N, block), leave=False, disable=not use_tqdm):
168 i1 = min(i0 + block, N)
169 # D2[i0:i1, :] = s[i0:i1,None] + s[None,:] - 2*X[i0:i1]@X.T
170 G = X[i0:i1] @ X.T # (B,N)
171 D2_blk = (s[i0:i1, None] + s[None, :] - 2.0 * G)
172 K = np.exp(-D2_blk * inv_two_sigma2)
173 if not cumulative:
174 row_sums[i0:i1] = np.sum(K, axis=1)
175 else:
176 K = np.cumsum(K, axis=1)
177 row_sums[i0:i1, 0:N] += np.triu(K, k=i0)
179 subvals = -np.log(np.clip(row_sums / N_mat, min=eps))/N_mat
180 entropy = np.sum(subvals, axis=0)
182 if not cumulative:
183 entropy = float(entropy)
185 return entropy, subvals
188def _get_entropy_torch(
189 descriptors: np.ndarray | torch.Tensor,
190 width: float,
191 use_tqdm: bool,
192 block: int,
193 eps: float,
194 cumulative: bool,
195 device: str | torch.device = 'cuda',
196 dtype: torch.dtype = torch.float64,
197) -> tuple[float, np.ndarray] | tuple[np.ndarray, np.ndarray]:
198 """Compute the informational entropy using torch.
199 See get_entropy for documentation.
200 """
201 with torch.no_grad():
202 # Move data to device
203 if not torch.cuda.is_available() and str(device) == 'cuda': # pragma: no cover
204 device = 'cpu'
205 X = torch.as_tensor(descriptors, dtype=dtype, device=device)
206 N, d = X.shape
207 inv_two_sigma2 = 1.0 / (2.0 * width * width)
209 # Precompute norms once
210 s = (X * X).sum(dim=1) # (N,)
212 if not cumulative:
213 N_mat = N
214 row_sums = torch.zeros(N, dtype=dtype, device=device)
215 else:
216 N_mat = torch.arange(1, N+1, dtype=dtype, device=device)
217 N_mat = N_mat.repeat(N).reshape((N, N))
218 row_sums = torch.tril(N_mat, diagonal=-1)
220 # Block over rows i; each block computes K[i0:i1, :].sum(-1)
221 # D2 = s[i] + s[j] - 2 * X[i] @ X[j]^T (formed in blocks to save memory)
222 XT = X.T # reuse in matmuls
223 for i0 in tqdm(range(0, N, block), leave=False, disable=not use_tqdm):
224 i1 = min(i0 + block, N)
225 try:
226 G = X[i0:i1] @ XT # (B, N)
227 D2_blk = s[i0:i1, None] + s[None, :] - 2.0 * G
228 except torch.cuda.OutOfMemoryError: # pragma: no cover
229 torch.cuda.empty_cache()
230 raise ValueError(
231 'Tried to allocate too much GPU memory.'
232 f' Try to reduce the value of block, e.g., to {block//2}.')
233 K = torch.exp(-D2_blk * inv_two_sigma2)
234 if not cumulative:
235 row_sums[i0:i1] = K.sum(dim=1)
236 else:
237 K = torch.cumsum(K, dim=1)
238 row_sums[i0:i1, 0:N] += torch.triu(K, diagonal=i0)
240 subvals = -torch.log(torch.clamp(row_sums / N_mat, min=eps))/N_mat
241 entropy = torch.sum(subvals, dim=0)
242 if not cumulative:
243 det_entropy = float(entropy.detach().cpu().item())
244 else:
245 det_entropy = entropy.detach().cpu().numpy()
246 det_subvals = subvals.detach().cpu().numpy()
247 return det_entropy, det_subvals