Coverage for calorine/tools/entropy.py: 100%
50 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-20 12:52 +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.float32,
55 eps: float = 1e-12,
56) -> tuple[float, np.ndarray]:
57 r"""
58 Computes an estimate for the information entropy :math:`H(\mathbf{X})` of a
59 set of descriptors :math:`\mathbf{X}`. The estimate is described in
60 [Nat. Comm. **16**, 4014 (2025)](https://doi.org/10.1038/s41467-025-59232-0)
61 and given by
63 .. math::
65 \mathcal{H}(\{\mathbf{X}\}) = -\frac{1}{n} \sum_{i=1}^{n} p_i
67 where
69 .. math::
71 p_i
72 = \log \left[
73 \frac{1}{n} \sum_{j=1}^{n}
74 K_h(\mathbf{X}_i, \mathbf{X}_j)
75 \right]
77 with a Gaussian kernel
79 .. math::
81 K_h(\mathbf{X}_i, \mathbf{X}_j)
82 = \exp\!\left(
83 -\frac{\lVert \mathbf{X}_i - \mathbf{X}_j \rVert^2}{2h^2}.
84 \right)
86 The calculation is done via torch if the library has been installed,
87 and numpy otherwise. When using torch the calculation is run via CUDA.
88 The latter behavior can be controlled using the :attr:`device` argument.
90 Parameters
91 ----------
92 descriptors
93 The set of descriptors :math:`\mathbf{X}` for which to evaluate to
94 the entropy. Typically each row corresponds to one atom and the
95 columns correspond to the different descriptor components.
96 width
97 Width :math:`h` of the Gaussian kernel.
98 use_tqdm
99 Use `tqdm <https://tqdm.github.io/>`_ to show a progress bar.
100 Note that this requires tqdm to be installed.
101 block
102 In order to limit the memory needs, the kernel density estimate
103 matrix is handled in blocks. This parameter controls the size of
104 each block. Smaller numbers imply a smaller memory footprint.
105 eps
106 Smallest (absolute) permissible value.
107 device
108 Device to use for calculation. The documentation of
109 [`torch.device`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device)
110 provides more information.
111 Only used when pytorch is available.
112 dtype
113 Floating point precision used for the computation. The documentation of
114 [`torch.dtype`](https://docs.pytorch.org/docs/stable/tensor_attributes.html)
115 provides an overview.
116 Only used when pytorch is available.
118 Returns
119 -------
120 A tuple comprising the total entropy :math:`H(\mathbf{X})` and the entropy contributions
121 :math:`p_i` from each row in the input descriptor matrix.
122 """
123 if torch_available:
124 res = _get_entropy_torch(descriptors, width, use_tqdm, block, eps, device, dtype)
125 return res
126 else: # pragma: no cover
127 warn('Using the numpy implementation.'
128 ' Install torch in order to use GPUs and achieve a considerable speed-up.')
129 res = _get_entropy_numpy(descriptors, width, use_tqdm, block, eps)
130 return res
133def _get_entropy_numpy(
134 descriptors: np.ndarray,
135 width: float,
136 use_tqdm: bool,
137 block: int,
138 eps: float,
139) -> float:
140 """Compute the informational entropy using numpy.
141 See get_entropy for documentation.
142 """
143 X = np.asarray(descriptors, dtype=np.float64)
144 N, d = X.shape
145 s = np.sum(X * X, axis=1) # (N,)
146 inv_two_sigma2 = 1.0 / (2.0 * width * width)
148 row_sums = np.zeros(N, dtype=np.float64)
149 for i0 in tqdm(range(0, N, block), leave=False, disable=not use_tqdm):
150 i1 = min(i0 + block, N)
151 # D2[i0:i1, :] = s[i0:i1,None] + s[None,:] - 2*X[i0:i1]@X.T
152 G = X[i0:i1] @ X.T # (B,N)
153 D2_blk = (s[i0:i1, None] + s[None, :] - 2.0 * G)
154 row_sums[i0:i1] = np.sum(np.exp(-D2_blk * inv_two_sigma2), axis=1)
156 subvals = -np.log(np.clip(row_sums / N, min=eps))
157 subvals /= N
158 entropy = np.sum(subvals)
159 return -float(entropy), subvals
162def _get_entropy_torch(
163 descriptors: np.ndarray | torch.Tensor,
164 width: float,
165 use_tqdm: bool,
166 block: int,
167 eps: float,
168 device: str | torch.device = 'cuda',
169 dtype: torch.dtype = torch.float32,
170) -> tuple[float, np.ndarray]:
171 """Compute the informational entropy using torch.
172 See get_entropy for documentation.
173 """
174 with torch.no_grad():
175 # Move data to device
176 if not torch.cuda.is_available() and str(device) == 'cuda': # pragma: no cover
177 device = 'cpu'
178 X = torch.as_tensor(descriptors, dtype=dtype, device=device)
179 N, d = X.shape
180 inv_two_sigma2 = 1.0 / (2.0 * width * width)
182 # Precompute norms once
183 s = (X * X).sum(dim=1) # (N,)
184 row_sums = torch.zeros(N, dtype=dtype, device=device)
186 # Block over rows i; each block computes K[i0:i1, :].sum(-1)
187 # D2 = s[i] + s[j] - 2 * X[i] @ X[j]^T (formed in blocks to save memory)
188 XT = X.T # reuse in matmuls
189 for i0 in tqdm(range(0, N, block), leave=False, disable=not use_tqdm):
190 i1 = min(i0 + block, N)
191 try:
192 G = X[i0:i1] @ XT # (B, N)
193 D2_blk = s[i0:i1, None] + s[None, :] - 2.0 * G
194 except torch.cuda.OutOfMemoryError: # pragma: no cover
195 torch.cuda.empty_cache()
196 raise ValueError(
197 'Tried to allocate too much GPU memory.'
198 f' Try to reduce the value of block, e.g., to {block//2}.')
199 row_sums[i0:i1] = torch.exp(-D2_blk * inv_two_sigma2).sum(dim=1)
201 subvals = -torch.log(torch.clamp(row_sums / N, min=eps))
202 subvals /= N
203 entropy = torch.sum(subvals)
204 return -float(entropy.detach().cpu().item()), subvals.detach().cpu().numpy()