Coverage for cosmolayer/cosmolayer.py: 97%

1""" 

2.. module:: cosmolayer.cosmolayer 

3 :synopsis: Differentiable COSMO-type activity coefficient layer. 

4""" 

5 

6import warnings 

7from collections.abc import Sequence 

8from typing import cast 

9 

10import numpy as np 

11import torch 

12from numpy.typing import NDArray 

13 

14from .cosmosolver import CosmoSolver 

15 

16AREA_PER_CONTACT = 79.53 # Ų 

17COORDINATION_NUMBER = 10 

18 

19 

20class CosmoLayer(torch.nn.Module): 

21 r"""Differentiable COSMO-type activity coefficient layer. 

22 

23 This class assumes that the interaction energy matrix :math:`\mathbf{U}` can depend 

24 on the temperature :math:`T` through the following relationship: 

25 

26 .. math:: 

27 

28 \frac{\mathbf{U}}{RT} = \sum_{n=1}^N \frac{\mathbf{U}_n}{RT^{\alpha_n}}, 

29 

30 

31 where each :math:`\mathbf{U}_n` is a constant interaction energy matrix, and 

32 :math:`\alpha_n` is a constant exponent. 

33 

34 To instantiate the class, the user must provide a reference temperature 

35 :math:`T_{\rm ref}`, a tuple of exponents :math:`(\alpha_1, \ldots, \alpha_N)`, and 

36 a tuple of scaled interaction energy matrices 

37 :math:`(\hat{\mathbf{U}}_1, \ldots, \hat{\mathbf{U}}_N)`, defined as: 

38 

39 .. math:: 

40 

41 \hat{\mathbf{U}}_n = \frac{\mathbf{U}_n}{RT_{\rm ref}^{\alpha_n}} 

42 

43 Parameters 

44 ---------- 

45 interaction_matrices : Sequence[NDArray[np.float64]] 

46 The scaled interaction energy matrices at the reference temperature 

47 (:math:`\hat{\mathbf{U}}_1, \ldots, \hat{\mathbf{U}}_N`). 

48 Must be square matrices, all with the same shape. 

49 exponents : Sequence[int] 

50 The temperature exponents :math:`(\alpha_1, \ldots, \alpha_N)`. Must have the 

51 same length as the number of interaction energy matrices. 

52 area_per_segment : float 

53 Area of each surface segment. 

54 reference_temperature : float, optional 

55 Reference temperature :math:`T_{\rm ref}`. Default is 298.15 K. 

56 learn_matrices : bool, optional 

57 Whether to register all scaled interaction energy matrices as trainable 

58 parameters. Default is False. 

59 implicit_diff : bool, optional 

60 Whether to differentiate through the Newton solver using the implicit 

61 function theorem (:class:`~cosmolayer.cosmosolver.CosmoSolver`), rather than 

62 backpropagating through the Newton iterations directly. Implicit 

63 differentiation is exact regardless of ``max_iter`` and is cheaper for 

64 backward passes, but requires the forward solve to have converged. Default 

65 is True. 

66 max_iter : int, optional 

67 Maximum number of iterations used by the Newton solver. If not 

68 given, defaults to 100 when ``implicit_diff`` is True and to 5 when it is 

69 False, since autograd-through-iterations becomes increasingly expensive to 

70 differentiate as ``max_iter`` grows. 

71 

72 Examples 

73 -------- 

74 >>> from importlib.resources import files 

75 >>> from cosmolayer import CosmoLayer 

76 >>> from cosmolayer.cosmosac import CosmoSac2002Model 

77 >>> import torch 

78 >>> T_ref = 298.15 # K 

79 >>> source = files("cosmolayer.data") 

80 >>> components = { 

81 ... "fluoromethane": (source / "CF.cosmo").read_text(), 

82 ... "water": (source / "O.cosmo").read_text(), 

83 ... } 

84 >>> mixture = CosmoSac2002Model.create_mixture(components) 

85 >>> interaction_matrices = mixture.interaction_matrices(T_ref) 

86 >>> exponents = mixture.temperature_exponents 

87 >>> area_per_segment = mixture.area_per_segment 

88 >>> cosmo_layer = CosmoLayer(interaction_matrices, exponents, area_per_segment) 

89 >>> cosmo_layer 

90 CosmoLayer( 

91 reference_temperature=298.15 

92 area_per_segment=7.50 

93 exponents=(1,) 

94 num_segment_types=51 

95 ) 

96 >>> T = torch.tensor(373.15) 

97 >>> x = torch.tensor([0.5, 0.5], requires_grad=True) 

98 >>> a = torch.tensor(mixture.areas) 

99 >>> v = torch.tensor(mixture.volumes) 

100 >>> P = torch.tensor(mixture.probabilities) 

101 >>> ln_gamma = cosmo_layer(T, x, a, v, P) 

102 >>> ln_gamma.tolist() 

103 [0.805809..., 0.648071...] 

104 >>> gE_RT = (x * ln_gamma).sum() 

105 >>> gE_RT.item() 

106 0.726940... 

107 >>> gE_RT.backward() 

108 >>> x.grad.tolist() 

109 [0.805809..., 0.648071...] 

110 """ 

111 

112 def __init__( # noqa: PLR0913 

113 self, 

114 interaction_matrices: Sequence[NDArray[np.float64]], 

115 exponents: Sequence[int], 

116 area_per_segment: float, 

117 *, 

118 reference_temperature: float = 298.15, # K 

119 learn_matrices: bool = False, 

120 implicit_diff: bool = True, 

121 max_iter: int | None = None, 

122 ): 

123 super().__init__() 

124 self._implicit_diff = implicit_diff 

125 if max_iter is None: 

126 max_iter = 100 if implicit_diff else 5 

127 self._max_iter = max_iter 

128 

129 num_matrices = len(interaction_matrices) 

130 if len(exponents) != num_matrices: 

131 raise ValueError( 

132 f"Number of exponents ({len(exponents)}) must match " 

133 f"number of interaction matrices ({num_matrices})" 

134 ) 

135 

136 self._num_matrices = num_matrices 

137 

138 shapes = {matrix.shape for matrix in interaction_matrices} 

139 if len(shapes) != 1: 

140 raise ValueError("All interaction matrices must have the same shape") 

141 rows, cols = shapes.pop() 

142 if rows != cols: 

143 raise ValueError("Interaction matrices must be square") 

144 self._n_types = rows 

145 

146 self._matrix_names_and_exponents: list[tuple[str, int]] = [] 

147 for idx, input_matrix in enumerate(interaction_matrices): 

148 matrix = torch.as_tensor(input_matrix) 

149 name = f"interaction_matrix_{idx}" 

150 if learn_matrices: 

151 self.register_parameter(name, torch.nn.Parameter(matrix)) 

152 else: 

153 self.register_buffer(name, matrix) 

154 self._matrix_names_and_exponents.append((name, exponents[idx])) 

155 

156 self._exponents = tuple(exponents) 

157 self._ref_temp = reference_temperature 

158 self._area_per_segment = area_per_segment 

159 self._kappa = COORDINATION_NUMBER / (2 * AREA_PER_CONTACT) 

160 self._max_iter = max_iter 

161 

162 def _check_convergence(self, converged: torch.Tensor) -> None: 

163 if not bool(converged.all()): 

164 warnings.warn( 

165 f"COSMO solver did not converge in {self._max_iter} iterations", 

166 RuntimeWarning, 

167 stacklevel=2, 

168 ) 

169 

170 def extra_repr(self) -> str: 

171 """Return a string representation of the CosmoLayer.""" 

172 return ( 

173 f"reference_temperature={self._ref_temp:.2f}\n" 

174 f"area_per_segment={self._area_per_segment:.2f}\n" 

175 f"exponents={self._exponents}\n" 

176 f"num_segment_types={self._n_types}" 

177 ) 

178 

179 def log_combinatorial_activity_coefficients( 

180 self, 

181 fracs: torch.Tensor, 

182 areas: torch.Tensor, 

183 volumes: torch.Tensor, 

184 ) -> torch.Tensor: 

185 r"""Compute the logarithms of the combinatorial activity coefficients. 

186 

187 Parameters 

188 ---------- 

189 fracs : torch.Tensor 

190 Mole fractions of the mixture components. 

191 Must sum to 1. Shape: (..., num_components). 

192 areas : torch.Tensor 

193 Surface areas of the mixture components, all in the same units. 

194 Shape: (..., num_components). 

195 volumes : torch.Tensor 

196 Volumes of the mixture components, all in the same units. 

197 Shape: (..., num_components). 

198 

199 Returns 

200 ------- 

201 torch.Tensor 

202 Logarithms of the combinatorial activity coefficients. 

203 Shape: (..., num_components). 

204 """ 

205 v_hat = volumes / (fracs * volumes).sum(dim=-1, keepdim=True) 

206 a_hat = areas / (fracs * areas).sum(dim=-1, keepdim=True) 

207 w_hat = v_hat / a_hat 

208 ln_gamma_c: torch.Tensor = ( 

209 1 - v_hat + v_hat.log() - self._kappa * areas * (1 - w_hat + w_hat.log()) 

210 ) 

211 return ln_gamma_c 

212 

213 def mixture_probabilities( 

214 self, 

215 fracs: torch.Tensor, 

216 areas: torch.Tensor, 

217 probs: torch.Tensor, 

218 ) -> torch.Tensor: 

219 """Compute the probabilities of segment types in the mixture. 

220 

221 Parameters 

222 ---------- 

223 fracs : torch.Tensor 

224 Mole fractions of the components. Must sum to 1. 

225 Shape: (..., num_components). 

226 areas : torch.Tensor 

227 Surface areas of the components. 

228 Shape: (..., num_components). 

229 probs : torch.Tensor 

230 Probabilities of segment types per component. Must sum to 1 along the 

231 segment type dimension. 

232 Shape: (..., num_components, num_segment_types). 

233 

234 Returns 

235 ------- 

236 torch.Tensor 

237 Probabilities of segment types in the mixture. 

238 Shape: (..., num_segment_types). 

239 """ 

240 xa = fracs * areas 

241 theta = xa / xa.sum(dim=-1, keepdim=True) 

242 return (theta.unsqueeze(-1) * probs).sum(dim=-2) 

243 

244 def scaled_interactions(self, temp: torch.Tensor) -> torch.Tensor: 

245 """Compute the scaled interactions at a given temperature. 

246 

247 Parameters 

248 ---------- 

249 temp : torch.Tensor 

250 Temperature in the same units as the reference temperature. 

251 Shape: (...,). 

252 

253 Returns 

254 ------- 

255 torch.Tensor 

256 The scaled interactions at the given temperature. 

257 Shape: (..., num_segment_types, num_segment_types). 

258 """ 

259 beta = (self._ref_temp / temp).unsqueeze(-1).unsqueeze(-1) 

260 matrices = [ 

261 getattr(self, name) * beta**exponent 

262 for name, exponent in self._matrix_names_and_exponents 

263 ] 

264 return torch.stack(matrices).sum(dim=0) 

265 

266 def log_pure_segment_activity_coefficients( 

267 self, 

268 scaled_interactions: torch.Tensor, 

269 probs: torch.Tensor, 

270 ) -> torch.Tensor: 

271 """Compute the log-activity coefficients of segment types in pure compounds. 

272 

273 Parameters 

274 ---------- 

275 scaled_interactions : torch.Tensor 

276 Scaled interaction energy matrix. 

277 Shape: (..., num_segment_types, num_segment_types). 

278 probs : torch.Tensor 

279 Probabilities of segment types per component. Must sum to 1 along the 

280 segment type dimension. 

281 Shape: (..., num_components, num_segment_types). 

282 

283 Returns 

284 ------- 

285 torch.Tensor 

286 Log-activity coefficients of segment types in pure compounds. 

287 Shape: (..., num_components, num_segment_types). 

288 

289 .. note:: 

290 Uses implicit differentiation through 

291 :meth:`CosmoSolver.apply <cosmolayer.cosmosolver.CosmoSolver.forward>` if 

292 ``implicit_diff`` is True, or backpropagates directly through the Newton 

293 iterations of :meth:`CosmoSolver.logspace_newton_solver 

294 <cosmolayer.cosmosolver.CosmoSolver.logspace_newton_solver>` otherwise. 

295 """ 

296 args = (probs, scaled_interactions.unsqueeze(-3), self._max_iter) 

297 if self._implicit_diff: 

298 log_gamma_pure, converged = CosmoSolver.apply(*args) 

299 else: 

300 log_gamma_pure, converged = CosmoSolver.logspace_newton_solver(*args) 

301 log_gamma_pure = log_gamma_pure.squeeze(-1) 

302 self._check_convergence(converged) 

303 return cast(torch.Tensor, log_gamma_pure) 

304 

305 def log_mixture_segment_activity_coefficients( 

306 self, 

307 scaled_interactions: torch.Tensor, 

308 fracs: torch.Tensor, 

309 areas: torch.Tensor, 

310 probs: torch.Tensor, 

311 ) -> torch.Tensor: 

312 """Compute the log-activity coefficients of segment types in the mixture. 

313 

314 Parameters 

315 ---------- 

316 scaled_interactions : torch.Tensor 

317 Scaled interaction energy matrix. 

318 Shape: (..., num_segment_types, num_segment_types). 

319 fracs : torch.Tensor 

320 Mole fractions of the components. Must sum to 1. 

321 Shape: (..., num_components). 

322 areas : torch.Tensor 

323 Surface areas of the components. 

324 Shape: (..., num_components). 

325 probs : torch.Tensor 

326 Probabilities of segment types per component. Must sum to 1 along the 

327 segment type dimension. 

328 Shape: (..., num_components, num_segment_types). 

329 

330 Returns 

331 ------- 

332 torch.Tensor 

333 Log-activity coefficients of segment types in the mixture. 

334 Shape: (..., num_segment_types). 

335 

336 .. note:: 

337 Uses implicit differentiation through 

338 :meth:`CosmoSolver.apply <cosmolayer.cosmosolver.CosmoSolver.forward>` if 

339 ``implicit_diff`` is True, or backpropagates directly through the Newton 

340 iterations of :meth:`CosmoSolver.logspace_newton_solver 

341 <cosmolayer.cosmosolver.CosmoSolver.logspace_newton_solver>` otherwise. 

342 """ 

343 args = ( 

344 self.mixture_probabilities(fracs, areas, probs), 

345 scaled_interactions, 

346 self._max_iter, 

347 ) 

348 if self._implicit_diff: 

349 log_gamma_mix, converged = CosmoSolver.apply(*args) 

350 else: 

351 log_gamma_mix, converged = CosmoSolver.logspace_newton_solver(*args) 

352 log_gamma_mix = log_gamma_mix.squeeze(-1) 

353 self._check_convergence(converged) 

354 return cast(torch.Tensor, log_gamma_mix) 

355 

356 def log_residual_activity_coefficients( 

357 self, 

358 temperature: torch.Tensor, 

359 fracs: torch.Tensor, 

360 areas: torch.Tensor, 

361 probs: torch.Tensor, 

362 ) -> torch.Tensor: 

363 """Compute the logarithms of the residual activity coefficients. 

364 

365 Parameters 

366 ---------- 

367 temperature : torch.Tensor 

368 Temperature in the same units as the reference temperature. 

369 Shape: (...,). 

370 fracs : torch.Tensor 

371 Mole fractions of the components. Must sum to 1. 

372 Shape: (..., num_components). 

373 areas : torch.Tensor 

374 Surface areas of the components. 

375 Shape: (..., num_components). 

376 probs : torch.Tensor 

377 Probabilities of segment types per component. Must sum to 1 along the 

378 segment type dimension. 

379 Shape: (..., num_components, num_segment_types). 

380 

381 Returns 

382 ------- 

383 torch.Tensor 

384 Logarithms of the residual activity coefficients. 

385 Shape: (..., num_components). 

386 """ 

387 scaled_interactions = self.scaled_interactions(temperature) 

388 log_gamma_pure = self.log_pure_segment_activity_coefficients( 

389 scaled_interactions, probs 

390 ) 

391 log_gamma_mix = self.log_mixture_segment_activity_coefficients( 

392 scaled_interactions, fracs, areas, probs 

393 ) 

394 num_segments = areas / self._area_per_segment 

395 log_gamma_res: torch.Tensor = num_segments * ( 

396 probs * (log_gamma_mix.unsqueeze(-2) - log_gamma_pure) 

397 ).sum(dim=-1) 

398 return log_gamma_res 

399 

400 def log_activity_coefficients( 

401 self, 

402 temperature: torch.Tensor, 

403 fracs: torch.Tensor, 

404 areas: torch.Tensor, 

405 volumes: torch.Tensor, 

406 probs: torch.Tensor, 

407 ) -> torch.Tensor: 

408 """Compute the logarithms of the activity coefficients. 

409 

410 Parameters 

411 ---------- 

412 temperature : torch.Tensor 

413 Temperature in the same units as the reference temperature. 

414 Shape: (...,). 

415 fracs : torch.Tensor 

416 Mole fractions of the components. Must sum to 1. 

417 Shape: (..., num_components). 

418 areas : torch.Tensor 

419 Surface areas of the components. 

420 Shape: (..., num_components). 

421 volumes : torch.Tensor 

422 Volumes of the components. 

423 Shape: (..., num_components). 

424 probs : torch.Tensor 

425 Probabilities of segment types per component. Must sum to 1 along the 

426 segment type dimension. 

427 Shape: (..., num_components, num_segment_types). 

428 

429 Returns 

430 ------- 

431 torch.Tensor 

432 Logarithms of the activity coefficients. 

433 Shape: (..., num_components). 

434 """ 

435 log_gamma_c = self.log_combinatorial_activity_coefficients( 

436 fracs, areas, volumes 

437 ) 

438 log_gamma_r = self.log_residual_activity_coefficients( 

439 temperature, fracs, areas, probs 

440 ) 

441 return log_gamma_c + log_gamma_r 

442 

443 def forward( 

444 self, 

445 temp: torch.Tensor, 

446 fracs: torch.Tensor, 

447 areas: torch.Tensor, 

448 volumes: torch.Tensor, 

449 probs: torch.Tensor, 

450 ) -> torch.Tensor: 

451 """Forward pass of the CosmoLayer. 

452 

453 Parameters 

454 ---------- 

455 temp : torch.Tensor 

456 Temperature in the same units as the reference temperature. 

457 Shape: (...,). 

458 fracs : torch.Tensor 

459 Mole fractions of the components. Must sum to 1. 

460 Shape: (..., num_components). 

461 areas : torch.Tensor 

462 Surface areas of the components. 

463 Shape: (..., num_components). 

464 volumes : torch.Tensor 

465 Volumes of the components. 

466 Shape: (..., num_components). 

467 probs : torch.Tensor 

468 Probabilities of segment types per component. Must sum to 1 along the 

469 segment type dimension. 

470 Shape: (..., num_components, num_segment_types). 

471 

472 Returns 

473 ------- 

474 torch.Tensor 

475 Logarithms of the activity coefficients. 

476 Shape: (..., num_components). 

477 """ 

478 return self.log_activity_coefficients(temp, fracs, areas, volumes, probs)