Coverage for cosmolayer / cosmolayer.py: 93%
74 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-11 14:25 +0000
1"""
2.. module:: cosmolayer.cosmolayer
3 :synopsis: Differentiable COSMO-type activity coefficient layer.
4"""
6from collections.abc import Sequence
7from typing import cast
9import numpy as np
10import torch
11from numpy.typing import NDArray
13from .cosmosolver import CosmoSolver
15AREA_PER_CONTACT = 79.53 # Ų
16COORDINATION_NUMBER = 10
19class CosmoLayer(torch.nn.Module):
20 r"""Differentiable COSMO-type activity coefficient layer.
22 This class assumes that the interaction energy matrix :math:`\mathbf{U}` can depend
23 on the temperature :math:`T` through the following relationship:
25 .. math::
27 \frac{\mathbf{U}}{RT} = \sum_{n=1}^N \frac{\mathbf{U}_n}{RT^{\alpha_n}},
30 where each :math:`\mathbf{U}_n` is a constant interaction energy matrix, and
31 :math:`\alpha_n` is a constant exponent.
33 To instantiate the class, the user must provide a reference temperature
34 :math:`T_{\rm ref}`, a tuple of exponents :math:`(\alpha_1, \ldots, \alpha_N)`, and
35 a tuple of scaled interaction energy matrices
36 :math:`(\hat{\mathbf{U}}_1, \ldots, \hat{\mathbf{U}}_N)`, defined as:
38 .. math::
40 \hat{\mathbf{U}}_n = \frac{\mathbf{U}_n}{RT_{\rm ref}^{\alpha_n}}
42 Parameters
43 ----------
44 interaction_matrices : Sequence[NDArray[np.float64]]
45 The scaled interaction energy matrices at the reference temperature
46 (:math:`\hat{\mathbf{U}}_1, \ldots, \hat{\mathbf{U}}_N`).
47 Must be square matrices, all with the same shape.
48 exponents : Sequence[int]
49 The temperature exponents :math:`(\alpha_1, \ldots, \alpha_N)`. Must have the
50 same length as the number of interaction energy matrices.
51 area_per_segment : float
52 Area of each surface segment.
53 reference_temperature : float, optional
54 Reference temperature :math:`T_{\rm ref}`. Default is 298.15 K.
55 learn_matrices : bool, optional
56 Whether to register all scaled interaction energy matrices as trainable
57 parameters. Default is False.
59 Examples
60 --------
61 >>> from importlib.resources import files
62 >>> from cosmolayer import CosmoLayer
63 >>> from cosmolayer.cosmosac import CosmoSac2002Model
64 >>> import torch
65 >>> T_ref = 298.15 # K
66 >>> source = files("cosmolayer.data")
67 >>> components = {
68 ... "fluoromethane": (source / "CF.cosmo").read_text(),
69 ... "water": (source / "O.cosmo").read_text(),
70 ... }
71 >>> mixture = CosmoSac2002Model.create_mixture(components)
72 >>> interaction_matrices = mixture.interaction_matrices(T_ref)
73 >>> exponents = mixture.temperature_exponents
74 >>> area_per_segment = mixture.area_per_segment
75 >>> cosmo_layer = CosmoLayer(interaction_matrices, exponents, area_per_segment)
76 >>> cosmo_layer
77 CosmoLayer(
78 reference_temperature=298.15
79 area_per_segment=7.50
80 exponents=(1,)
81 num_segment_types=51
82 )
83 >>> T = torch.tensor(373.15)
84 >>> x = torch.tensor([0.5, 0.5], requires_grad=True)
85 >>> a = torch.tensor(mixture.areas)
86 >>> v = torch.tensor(mixture.volumes)
87 >>> P = torch.tensor(mixture.probabilities)
88 >>> ln_gamma = cosmo_layer(T, x, a, v, P)
89 >>> ln_gamma.tolist()
90 [0.805809..., 0.648071...]
91 >>> gE_RT = (x * ln_gamma).sum()
92 >>> gE_RT.item()
93 0.726940...
94 >>> gE_RT.backward()
95 >>> x.grad.tolist()
96 [0.805809..., 0.648071...]
97 """
99 def __init__( # noqa: PLR0913
100 self,
101 interaction_matrices: Sequence[NDArray[np.float64]],
102 exponents: Sequence[int],
103 area_per_segment: float,
104 *,
105 reference_temperature: float = 298.15, # K
106 learn_matrices: bool = False,
107 max_iter: int = 100,
108 ):
109 super().__init__()
111 num_matrices = len(interaction_matrices)
112 if len(exponents) != num_matrices:
113 raise ValueError(
114 f"Number of exponents ({len(exponents)}) must match "
115 f"number of interaction matrices ({num_matrices})"
116 )
118 self._num_matrices = num_matrices
120 shapes = {matrix.shape for matrix in interaction_matrices}
121 if len(shapes) != 1:
122 raise ValueError("All interaction matrices must have the same shape")
123 rows, cols = shapes.pop()
124 if rows != cols:
125 raise ValueError("Interaction matrices must be square")
126 self._n_types = rows
128 self._matrix_names_and_exponents: list[tuple[str, int]] = []
129 for idx, input_matrix in enumerate(interaction_matrices):
130 matrix = torch.as_tensor(input_matrix)
131 name = f"interaction_matrix_{idx}"
132 if learn_matrices:
133 self.register_parameter(name, torch.nn.Parameter(matrix))
134 else:
135 self.register_buffer(name, matrix)
136 self._matrix_names_and_exponents.append((name, exponents[idx]))
138 self._exponents = tuple(exponents)
139 self._ref_temp = reference_temperature
140 self._area_per_segment = area_per_segment
141 self._kappa = COORDINATION_NUMBER / (2 * AREA_PER_CONTACT)
142 self._max_iter = max_iter
144 def extra_repr(self) -> str:
145 """Return a string representation of the CosmoLayer."""
146 return (
147 f"reference_temperature={self._ref_temp:.2f}\n"
148 f"area_per_segment={self._area_per_segment:.2f}\n"
149 f"exponents={self._exponents}\n"
150 f"num_segment_types={self._n_types}"
151 )
153 def log_combinatorial_activity_coefficients(
154 self,
155 fracs: torch.Tensor,
156 areas: torch.Tensor,
157 volumes: torch.Tensor,
158 ) -> torch.Tensor:
159 r"""Compute the logarithms of the combinatorial activity coefficients.
161 Parameters
162 ----------
163 fracs : torch.Tensor
164 Mole fractions of the mixture components.
165 Must sum to 1. Shape: (..., num_components).
166 areas : torch.Tensor
167 Surface areas of the mixture components, all in the same units.
168 Shape: (..., num_components).
169 volumes : torch.Tensor
170 Volumes of the mixture components, all in the same units.
171 Shape: (..., num_components).
173 Returns
174 -------
175 torch.Tensor
176 Logarithms of the combinatorial activity coefficients.
177 Shape: (..., num_components).
178 """
179 v_hat = volumes / (fracs * volumes).sum(dim=-1, keepdim=True)
180 a_hat = areas / (fracs * areas).sum(dim=-1, keepdim=True)
181 w_hat = v_hat / a_hat
182 ln_gamma_c: torch.Tensor = (
183 1 - v_hat + v_hat.log() - self._kappa * areas * (1 - w_hat + w_hat.log())
184 )
185 return ln_gamma_c
187 def mixture_probabilities(
188 self,
189 fracs: torch.Tensor,
190 areas: torch.Tensor,
191 probs: torch.Tensor,
192 ) -> torch.Tensor:
193 """Compute the probabilities of segment types in the mixture.
195 Parameters
196 ----------
197 fracs : torch.Tensor
198 Mole fractions of the components. Must sum to 1.
199 Shape: (..., num_components).
200 areas : torch.Tensor
201 Surface areas of the components.
202 Shape: (..., num_components).
203 probs : torch.Tensor
204 Probabilities of segment types per component. Must sum to 1 along the
205 segment type dimension.
206 Shape: (..., num_components, num_segment_types).
208 Returns
209 -------
210 torch.Tensor
211 Probabilities of segment types in the mixture.
212 Shape: (..., num_segment_types).
213 """
214 xa = fracs * areas
215 theta = xa / xa.sum(dim=-1, keepdim=True)
216 return (theta.unsqueeze(-1) * probs).sum(dim=-2)
218 def scaled_interactions(self, temp: torch.Tensor) -> torch.Tensor:
219 """Compute the scaled interactions at a given temperature.
221 Parameters
222 ----------
223 temp : torch.Tensor
224 Temperature in the same units as the reference temperature.
225 Shape: (...,).
227 Returns
228 -------
229 torch.Tensor
230 The scaled interactions at the given temperature.
231 Shape: (..., num_segment_types, num_segment_types).
232 """
233 beta = (self._ref_temp / temp).unsqueeze(-1).unsqueeze(-1)
234 matrices = [
235 getattr(self, name) * beta**exponent
236 for name, exponent in self._matrix_names_and_exponents
237 ]
238 return torch.stack(matrices).sum(dim=0)
240 def log_pure_segment_activity_coefficients(
241 self,
242 scaled_interactions: torch.Tensor,
243 probs: torch.Tensor,
244 ) -> torch.Tensor:
245 """Compute the log-activity coefficients of segment types in pure compounds.
247 Parameters
248 ----------
249 scaled_interactions : torch.Tensor
250 Scaled interaction energy matrix.
251 Shape: (..., num_segment_types, num_segment_types).
252 probs : torch.Tensor
253 Probabilities of segment types per component. Must sum to 1 along the
254 segment type dimension.
255 Shape: (..., num_components, num_segment_types).
257 Returns
258 -------
259 torch.Tensor
260 Log-activity coefficients of segment types in pure compounds.
261 Shape: (..., num_components, num_segment_types).
262 """
263 log_gamma_pure, converged = CosmoSolver.apply(
264 probs, scaled_interactions.unsqueeze(-3), self._max_iter
265 )
266 if not bool(converged.all()):
267 raise RuntimeError(
268 f"Newton solver did not converge in {self._max_iter} iterations"
269 )
270 return cast(torch.Tensor, log_gamma_pure)
272 def log_mixture_segment_activity_coefficients(
273 self,
274 scaled_interactions: torch.Tensor,
275 fracs: torch.Tensor,
276 areas: torch.Tensor,
277 probs: torch.Tensor,
278 ) -> torch.Tensor:
279 """Compute the log-activity coefficients of segment types in the mixture.
281 Parameters
282 ----------
283 scaled_interactions : torch.Tensor
284 Scaled interaction energy matrix.
285 Shape: (..., num_segment_types, num_segment_types).
286 fracs : torch.Tensor
287 Mole fractions of the components. Must sum to 1.
288 Shape: (..., num_components).
289 areas : torch.Tensor
290 Surface areas of the components.
291 Shape: (..., num_components).
292 probs : torch.Tensor
293 Probabilities of segment types per component. Must sum to 1 along the
294 segment type dimension.
295 Shape: (..., num_components, num_segment_types).
297 Returns
298 -------
299 torch.Tensor
300 Log-activity coefficients of segment types in the mixture.
301 Shape: (..., num_segment_types).
302 """
303 log_gamma_mix, converged = CosmoSolver.apply(
304 self.mixture_probabilities(fracs, areas, probs),
305 scaled_interactions,
306 self._max_iter,
307 )
308 if not bool(converged.all()):
309 raise RuntimeError(
310 f"Newton solver did not converge in {self._max_iter} iterations"
311 )
312 return cast(torch.Tensor, log_gamma_mix)
314 def log_residual_activity_coefficients(
315 self,
316 temperature: torch.Tensor,
317 fracs: torch.Tensor,
318 areas: torch.Tensor,
319 probs: torch.Tensor,
320 ) -> torch.Tensor:
321 """Compute the logarithms of the residual activity coefficients.
323 Parameters
324 ----------
325 temperature : torch.Tensor
326 Temperature in the same units as the reference temperature.
327 Shape: (...,).
328 fracs : torch.Tensor
329 Mole fractions of the components. Must sum to 1.
330 Shape: (..., num_components).
331 areas : torch.Tensor
332 Surface areas of the components.
333 Shape: (..., num_components).
334 probs : torch.Tensor
335 Probabilities of segment types per component. Must sum to 1 along the
336 segment type dimension.
337 Shape: (..., num_components, num_segment_types).
339 Returns
340 -------
341 torch.Tensor
342 Logarithms of the residual activity coefficients.
343 Shape: (..., num_components).
344 """
345 scaled_interactions = self.scaled_interactions(temperature)
346 log_gamma_pure = self.log_pure_segment_activity_coefficients(
347 scaled_interactions, probs
348 )
349 log_gamma_mix = self.log_mixture_segment_activity_coefficients(
350 scaled_interactions, fracs, areas, probs
351 )
352 num_segments = areas / self._area_per_segment
353 log_gamma_res: torch.Tensor = num_segments * (
354 probs * (log_gamma_mix.unsqueeze(-2) - log_gamma_pure)
355 ).sum(dim=-1)
356 return log_gamma_res
358 def log_activity_coefficients(
359 self,
360 temperature: torch.Tensor,
361 fracs: torch.Tensor,
362 areas: torch.Tensor,
363 volumes: torch.Tensor,
364 probs: torch.Tensor,
365 ) -> torch.Tensor:
366 """Compute the logarithms of the activity coefficients.
368 Parameters
369 ----------
370 temperature : torch.Tensor
371 Temperature in the same units as the reference temperature.
372 Shape: (...,).
373 fracs : torch.Tensor
374 Mole fractions of the components. Must sum to 1.
375 Shape: (..., num_components).
376 areas : torch.Tensor
377 Surface areas of the components.
378 Shape: (..., num_components).
379 volumes : torch.Tensor
380 Volumes of the components.
381 Shape: (..., num_components).
382 probs : torch.Tensor
383 Probabilities of segment types per component. Must sum to 1 along the
384 segment type dimension.
385 Shape: (..., num_components, num_segment_types).
387 Returns
388 -------
389 torch.Tensor
390 Logarithms of the activity coefficients.
391 Shape: (..., num_components).
392 """
393 log_gamma_c = self.log_combinatorial_activity_coefficients(
394 fracs, areas, volumes
395 )
396 log_gamma_r = self.log_residual_activity_coefficients(
397 temperature, fracs, areas, probs
398 )
399 return log_gamma_c + log_gamma_r
401 def forward(
402 self,
403 temp: torch.Tensor,
404 fracs: torch.Tensor,
405 areas: torch.Tensor,
406 volumes: torch.Tensor,
407 probs: torch.Tensor,
408 ) -> torch.Tensor:
409 """Forward pass of the CosmoLayer.
411 Parameters
412 ----------
413 temp : torch.Tensor
414 Temperature in the same units as the reference temperature.
415 Shape: (...,).
416 fracs : torch.Tensor
417 Mole fractions of the components. Must sum to 1.
418 Shape: (..., num_components).
419 areas : torch.Tensor
420 Surface areas of the components.
421 Shape: (..., num_components).
422 volumes : torch.Tensor
423 Volumes of the components.
424 Shape: (..., num_components).
425 probs : torch.Tensor
426 Probabilities of segment types per component. Must sum to 1 along the
427 segment type dimension.
428 Shape: (..., num_components, num_segment_types).
430 Returns
431 -------
432 torch.Tensor
433 Logarithms of the activity coefficients.
434 Shape: (..., num_components).
435 """
436 return self.log_activity_coefficients(temp, fracs, areas, volumes, probs)