Coverage for cosmolayer / cosmolayer.py: 95%

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 

60 Examples 

61 -------- 

62 >>> from importlib.resources import files 

63 >>> from cosmolayer import CosmoLayer 

64 >>> from cosmolayer.cosmosac import CosmoSac2002Model 

65 >>> import torch 

66 >>> T_ref = 298.15 # K 

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

68 >>> components = { 

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

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

71 ... } 

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

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

74 >>> exponents = mixture.temperature_exponents 

75 >>> area_per_segment = mixture.area_per_segment 

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

77 >>> cosmo_layer 

78 CosmoLayer( 

79 reference_temperature=298.15 

80 area_per_segment=7.50 

81 exponents=(1,) 

82 num_segment_types=51 

83 ) 

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

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

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

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

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

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

90 >>> ln_gamma.tolist() 

91 [0.805809..., 0.648071...] 

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

93 >>> gE_RT.item() 

94 0.726940... 

95 >>> gE_RT.backward() 

96 >>> x.grad.tolist() 

97 [0.805809..., 0.648071...] 

98 """ 

99 

100 def __init__( # noqa: PLR0913 

101 self, 

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

103 exponents: Sequence[int], 

104 area_per_segment: float, 

105 *, 

106 reference_temperature: float = 298.15, # K 

107 learn_matrices: bool = False, 

108 max_iter: int = 100, 

109 ): 

110 super().__init__() 

111 

112 num_matrices = len(interaction_matrices) 

113 if len(exponents) != num_matrices: 

114 raise ValueError( 

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

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

117 ) 

118 

119 self._num_matrices = num_matrices 

120 

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

122 if len(shapes) != 1: 

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

124 rows, cols = shapes.pop() 

125 if rows != cols: 

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

127 self._n_types = rows 

128 

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

130 for idx, input_matrix in enumerate(interaction_matrices): 

131 matrix = torch.as_tensor(input_matrix) 

132 name = f"interaction_matrix_{idx}" 

133 if learn_matrices: 

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

135 else: 

136 self.register_buffer(name, matrix) 

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

138 

139 self._exponents = tuple(exponents) 

140 self._ref_temp = reference_temperature 

141 self._area_per_segment = area_per_segment 

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

143 self._max_iter = max_iter 

144 

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

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

147 warnings.warn( 

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

149 RuntimeWarning, 

150 stacklevel=2, 

151 ) 

152 

153 def extra_repr(self) -> str: 

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

155 return ( 

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

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

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

159 f"num_segment_types={self._n_types}" 

160 ) 

161 

162 def log_combinatorial_activity_coefficients( 

163 self, 

164 fracs: torch.Tensor, 

165 areas: torch.Tensor, 

166 volumes: torch.Tensor, 

167 ) -> torch.Tensor: 

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

169 

170 Parameters 

171 ---------- 

172 fracs : torch.Tensor 

173 Mole fractions of the mixture components. 

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

175 areas : torch.Tensor 

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

177 Shape: (..., num_components). 

178 volumes : torch.Tensor 

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

180 Shape: (..., num_components). 

181 

182 Returns 

183 ------- 

184 torch.Tensor 

185 Logarithms of the combinatorial activity coefficients. 

186 Shape: (..., num_components). 

187 """ 

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

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

190 w_hat = v_hat / a_hat 

191 ln_gamma_c: torch.Tensor = ( 

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

193 ) 

194 return ln_gamma_c 

195 

196 def mixture_probabilities( 

197 self, 

198 fracs: torch.Tensor, 

199 areas: torch.Tensor, 

200 probs: torch.Tensor, 

201 ) -> torch.Tensor: 

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

203 

204 Parameters 

205 ---------- 

206 fracs : torch.Tensor 

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

208 Shape: (..., num_components). 

209 areas : torch.Tensor 

210 Surface areas of the components. 

211 Shape: (..., num_components). 

212 probs : torch.Tensor 

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

214 segment type dimension. 

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

216 

217 Returns 

218 ------- 

219 torch.Tensor 

220 Probabilities of segment types in the mixture. 

221 Shape: (..., num_segment_types). 

222 """ 

223 xa = fracs * areas 

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

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

226 

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

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

229 

230 Parameters 

231 ---------- 

232 temp : torch.Tensor 

233 Temperature in the same units as the reference temperature. 

234 Shape: (...,). 

235 

236 Returns 

237 ------- 

238 torch.Tensor 

239 The scaled interactions at the given temperature. 

240 Shape: (..., num_segment_types, num_segment_types). 

241 """ 

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

243 matrices = [ 

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

245 for name, exponent in self._matrix_names_and_exponents 

246 ] 

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

248 

249 def log_pure_segment_activity_coefficients( 

250 self, 

251 scaled_interactions: torch.Tensor, 

252 probs: torch.Tensor, 

253 ) -> torch.Tensor: 

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

255 

256 Parameters 

257 ---------- 

258 scaled_interactions : torch.Tensor 

259 Scaled interaction energy matrix. 

260 Shape: (..., num_segment_types, num_segment_types). 

261 probs : torch.Tensor 

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

263 segment type dimension. 

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

265 

266 Returns 

267 ------- 

268 torch.Tensor 

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

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

271 """ 

272 log_gamma_pure, converged = CosmoSolver.apply( 

273 probs, scaled_interactions.unsqueeze(-3), self._max_iter 

274 ) 

275 self._check_convergence(converged) 

276 return cast(torch.Tensor, log_gamma_pure) 

277 

278 def log_mixture_segment_activity_coefficients( 

279 self, 

280 scaled_interactions: torch.Tensor, 

281 fracs: torch.Tensor, 

282 areas: torch.Tensor, 

283 probs: torch.Tensor, 

284 ) -> torch.Tensor: 

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

286 

287 Parameters 

288 ---------- 

289 scaled_interactions : torch.Tensor 

290 Scaled interaction energy matrix. 

291 Shape: (..., num_segment_types, num_segment_types). 

292 fracs : torch.Tensor 

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

294 Shape: (..., num_components). 

295 areas : torch.Tensor 

296 Surface areas of the components. 

297 Shape: (..., num_components). 

298 probs : torch.Tensor 

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

300 segment type dimension. 

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

302 

303 Returns 

304 ------- 

305 torch.Tensor 

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

307 Shape: (..., num_segment_types). 

308 """ 

309 log_gamma_mix, converged = CosmoSolver.apply( 

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

311 scaled_interactions, 

312 self._max_iter, 

313 ) 

314 self._check_convergence(converged) 

315 return cast(torch.Tensor, log_gamma_mix) 

316 

317 def log_residual_activity_coefficients( 

318 self, 

319 temperature: torch.Tensor, 

320 fracs: torch.Tensor, 

321 areas: torch.Tensor, 

322 probs: torch.Tensor, 

323 ) -> torch.Tensor: 

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

325 

326 Parameters 

327 ---------- 

328 temperature : torch.Tensor 

329 Temperature in the same units as the reference temperature. 

330 Shape: (...,). 

331 fracs : torch.Tensor 

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

333 Shape: (..., num_components). 

334 areas : torch.Tensor 

335 Surface areas of the components. 

336 Shape: (..., num_components). 

337 probs : torch.Tensor 

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

339 segment type dimension. 

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

341 

342 Returns 

343 ------- 

344 torch.Tensor 

345 Logarithms of the residual activity coefficients. 

346 Shape: (..., num_components). 

347 """ 

348 scaled_interactions = self.scaled_interactions(temperature) 

349 log_gamma_pure = self.log_pure_segment_activity_coefficients( 

350 scaled_interactions, probs 

351 ) 

352 log_gamma_mix = self.log_mixture_segment_activity_coefficients( 

353 scaled_interactions, fracs, areas, probs 

354 ) 

355 num_segments = areas / self._area_per_segment 

356 log_gamma_res: torch.Tensor = num_segments * ( 

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

358 ).sum(dim=-1) 

359 return log_gamma_res 

360 

361 def log_activity_coefficients( 

362 self, 

363 temperature: torch.Tensor, 

364 fracs: torch.Tensor, 

365 areas: torch.Tensor, 

366 volumes: torch.Tensor, 

367 probs: torch.Tensor, 

368 ) -> torch.Tensor: 

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

370 

371 Parameters 

372 ---------- 

373 temperature : torch.Tensor 

374 Temperature in the same units as the reference temperature. 

375 Shape: (...,). 

376 fracs : torch.Tensor 

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

378 Shape: (..., num_components). 

379 areas : torch.Tensor 

380 Surface areas of the components. 

381 Shape: (..., num_components). 

382 volumes : torch.Tensor 

383 Volumes of the components. 

384 Shape: (..., num_components). 

385 probs : torch.Tensor 

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

387 segment type dimension. 

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

389 

390 Returns 

391 ------- 

392 torch.Tensor 

393 Logarithms of the activity coefficients. 

394 Shape: (..., num_components). 

395 """ 

396 log_gamma_c = self.log_combinatorial_activity_coefficients( 

397 fracs, areas, volumes 

398 ) 

399 log_gamma_r = self.log_residual_activity_coefficients( 

400 temperature, fracs, areas, probs 

401 ) 

402 return log_gamma_c + log_gamma_r 

403 

404 def forward( 

405 self, 

406 temp: torch.Tensor, 

407 fracs: torch.Tensor, 

408 areas: torch.Tensor, 

409 volumes: torch.Tensor, 

410 probs: torch.Tensor, 

411 ) -> torch.Tensor: 

412 """Forward pass of the CosmoLayer. 

413 

414 Parameters 

415 ---------- 

416 temp : torch.Tensor 

417 Temperature in the same units as the reference temperature. 

418 Shape: (...,). 

419 fracs : torch.Tensor 

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

421 Shape: (..., num_components). 

422 areas : torch.Tensor 

423 Surface areas of the components. 

424 Shape: (..., num_components). 

425 volumes : torch.Tensor 

426 Volumes of the components. 

427 Shape: (..., num_components). 

428 probs : torch.Tensor 

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

430 segment type dimension. 

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

432 

433 Returns 

434 ------- 

435 torch.Tensor 

436 Logarithms of the activity coefficients. 

437 Shape: (..., num_components). 

438 """ 

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