Coverage for cosmolayer/cosmosac/component.py: 94%

1import os

2import sys

3from typing import Any, TextIO

5if sys.version_info >= (3, 11):

6 from importlib.resources.abc import Traversable

7else:

8 Traversable = Any

10import numpy as np

11import pandas as pd

12import periodictable as pt

13from numpy.typing import NDArray

15from ..parser import parse_cosmo_file

16from .constants import (

17 COSMO_SAC_2010_AVERAGING_RADIUS,

18 COSMO_SAC_2010_F_DECAY,

19 COSMO_SAC_2010_SIGMA_0,

20)

21from .segment_groups import NHB, OH, OT, SEGMENT_GROUPS

23COVALENT_FACTOR = 1.3 # Same as in RDKit

26class Component:

27 r"""Molecular component for the COSMO-SAC activity coefficient model.

29 Parameters

30 ----------

31 cosmo_string : str

32 Contents of a COSMO output file from quantum mechanical calculations.

34 Keyword Arguments

35 -----------------

36 min_sigma : float, optional

37 Minimum screening charge density in e/Å². Default is -0.025 e/Å².

38 max_sigma : float, optional

39 Maximum screening charge density in e/Å². Default is 0.025 e/Å².

40 num_points : int, optional

41 Number of discrete points in the sigma profile. Default is 51.

42 averaging_radius : float, optional

43 Effective radius for distance-weighted sigma averaging in Å.

44 Default is √(7.25 / π) Å :cite:`Bell2020`.

45 f_decay : float, optional

46 Decay factor for exponential distance weighting in the sigma averaging

47 procedure. Default is 3.57 :cite:`Bell2020`.

48 sigma_0 : float or None, optional

49 Standard deviation of the Gaussian probability of a segment to form a hydrogen

50 bond in e/Å². Set to ``None`` to disable hydrogen-bond splitting (all

51 surface area is assigned to the NHB class).

52 Default is 0.007 e/Å² :cite:`Bell2020`.

53 merge_profiles : bool, optional

54 Whether to merge segment groups (NHB, OH, OT) into a single profile

55 when accessing :attr:`probabilities` and :attr:`sigma_profile`.

56 Default is False.

58 Raises

59 ------

60 ValueError

61 If the COSMO string is not in any supported format.

62 ValueError

63 If averaged charge densities fall outside the specified sigma range.

65 Examples

66 --------

67 >>> import numpy as np

68 >>> from importlib.resources import files

69 >>> from cosmolayer.cosmosac import Component

70 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

71 >>> component = Component(path.read_text())

72 >>> component.area

73 97.34554...

74 >>> component.volume

75 80.07160...

77 When :attr:`merge_profiles` is True, :attr:`sigma_profile` is a single

78 merged profile:

80 >>> component = Component(path.read_text(), merge_profiles=True)

81 >>> sigma_profile = component.sigma_profile

82 >>> sigma_profile.shape

83 (51,)

84 >>> print(sum(sigma_profile))

85 97.34554...

87 When :attr:`merge_profiles` is False, :attr:`sigma_profile` is stacked

88 (NHB, OH, OT), shape (3, num_points):

90 >>> component = Component(path.read_text(), merge_profiles=False)

91 >>> stacked = component.sigma_profile

92 >>> stacked.shape

93 (3, 51)

94 >>> from cosmolayer.cosmosac.segment_groups import SEGMENT_GROUPS

95 >>> for i, s in enumerate(SEGMENT_GROUPS):

96 ... print(s, sum(stacked[i]))

97 NHB 72.31802...

98 OH 12.25732...

99 OT 12.77019...

100

101 Plotting the sigma profiles (stacked, :attr:`merge_profiles` is False):

102

103

104 .. plot::

105 :context: close-figs

106

107 >>> from importlib.resources import files

108 >>> from cosmolayer.cosmosac import Component

109 >>> from cosmolayer.cosmosac.segment_groups import SEGMENT_GROUPS

110 >>> from matplotlib import pyplot as plt

111 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

112 >>> component = Component(path.read_text(), merge_profiles=False)

113 >>> fig, ax = plt.subplots(figsize=(8, 4))

114 >>> grid = component.sigma_grid

115 >>> for i, label in enumerate(SEGMENT_GROUPS):

116 ... _ = ax.plot(grid, component.sigma_profile[i], label=label)

117 >>> _ = ax.set_xlabel("Charge density (e/Å²)")

118 >>> _ = ax.set_ylabel("Surface area contribution (Å²)")

119 >>> _ = ax.legend()

120 >>> fig.tight_layout()

121

122 Plotting the segment-type probabilities:

123

124 .. plot::

125 :context: close-figs

126

127 >>> from importlib.resources import files

128 >>> from cosmolayer.cosmosac import Component

129 >>> from matplotlib import pyplot as plt

130 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

131 >>> component = Component(path.read_text())

132 >>> fig, ax = plt.subplots(figsize=(8, 4))

133 >>> p = component.probabilities

134 >>> _ = ax.bar(range(len(p)), p)

135 >>> _ = ax.set_xlabel("Segment type index")

136 >>> _ = ax.set_ylabel("Probability")

137 >>> fig.tight_layout()

138 """

139

140 def __init__( # noqa: PLR0913

141 self,

142 cosmo_string: str,

143 *,

144 min_sigma: float = -0.025, # e/Å²

145 max_sigma: float = 0.025, # e/Å²

146 num_points: int = 51,

147 averaging_radius: float = COSMO_SAC_2010_AVERAGING_RADIUS, # Å

148 f_decay: float = COSMO_SAC_2010_F_DECAY,

149 sigma_0: float | None = COSMO_SAC_2010_SIGMA_0, # e/Å²

150 merge_profiles: bool = False,

151 ) -> None:

152 self._min_sigma = min_sigma

153 self._grid = np.linspace(min_sigma, max_sigma, num_points)

154 self._bin_width = (max_sigma - min_sigma) / (num_points - 1)

155

156 self._averaging_radius = averaging_radius

157 self._f_decay = f_decay

158 self._sigma_0 = sigma_0

159 self._merge_profiles = merge_profiles

160

161 self._format, self._atom_data, self._segment_data, self._volume = (

162 parse_cosmo_file(cosmo_string)

163 )

164

165 sigmas, averaged_sigmas = self._average_sigmas()

166 if (averaged_sigmas < min_sigma).any() or (averaged_sigmas > max_sigma).any():

167 raise ValueError("Averaged charge densities out of range.")

168 self._segment_data["sigma"] = sigmas

169 self._segment_data["sigma_avg"] = averaged_sigmas

170

171 self._bonds = self._detect_bonds()

172 self._area = float(self._segment_data["area"].sum())

173 self._sigma_profiles = self._compute_sigma_profiles(averaged_sigmas)

174

175 def __repr__(self) -> str:

176 num_atoms = len(self._atom_data)

177 num_segments = len(self._segment_data)

178 return f"Component({num_atoms} atoms, {num_segments} segments)"

179

180 @staticmethod

181 def _get_covalent_radius(element: str) -> float:

182 """Get scaled covalent radius for bond detection.

183

184 Parameters

185 ----------

186 element : str

187 Chemical element symbol.

188

189 Returns

190 -------

191 float

192 Covalent radius in Å, scaled by factor 1.3.

193 """

194 covalent_radius = pt.elements.symbol(element).covalent_radius

195 if covalent_radius is None:

196 raise ValueError(f"Unknown covalent radius for element {element!r}")

197 return COVALENT_FACTOR * float(covalent_radius)

198

199 def _detect_bonds(self) -> list[tuple[int, int]]:

200 """Determines bonds from interatomic distances."""

201 df = self._atom_data

202 coords = df[["x", "y", "z"]].values

203 distances = np.sqrt(np.square(coords[:, None, :] - coords).sum(axis=-1))

204 radii = df["element"].apply(self._get_covalent_radius).values

205 adjacency_matrix = distances < (radii[:, None] + radii[None, :])

206 bond_indices = np.nonzero(np.triu(adjacency_matrix, k=1))

207 return [(int(i), int(j)) for i, j in zip(*bond_indices, strict=True)]

208

209 def _get_hydrogen_bonding_classes(self) -> pd.Series:

210 """Classify atoms into hydrogen bonding types (OH, OT, NHB).

211

212 Assigns hydrogen bonding classes: OH (O-H bonds), OT (N-H, F-H bonds or

213 isolated N/F/O), and NHB (all other atoms).

214

215 Returns

216 -------

217 pd.Series

218 Hydrogen bonding class label for each atom.

219 """

220 elements = self._atom_data["element"]

221 hb_class = elements.apply(

222 lambda element: OT if element in ["N", "F", "O"] else NHB

223 )

224 for i, j in self._bonds:

225 elements_ij = set(elements.iloc[[i, j]])

226 if elements_ij in [{"O", "H"}, {"N", "H"}, {"F", "H"}]:

227 hb_class.at[i] = hb_class.at[j] = OH if "O" in elements_ij else OT

228 return hb_class

229

230 def _average_sigmas(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]:

231 """Apply distance-weighted averaging to segment charge densities.

232

233 Smooths raw screening charge densities (σ = q/A) using exponentially

234 decaying weights based on distances between segment centroids.

235

236 Returns

237 -------

238 np.ndarray

239 Averaged screening charge density for each segment in e/Å.

240 """

241 sigmas = self._segment_data["charge"].values / self._segment_data["area"].values

242 coords = self._segment_data[["x", "y", "z"]].values

243 squared_distances = np.square(coords[:, None, :] - coords).sum(axis=-1)

244 squared_radii = self._segment_data["area"].values / np.pi

245

246 sums = squared_radii + self._averaging_radius**2

247 prods = squared_radii * self._averaging_radius**2

248 weights = np.exp(-self._f_decay * squared_distances / sums) * prods / sums

249 averaged_sigmas: NDArray[np.float64] = np.sum(

250 weights * sigmas, axis=1

251 ) / np.sum(weights, axis=1)

252

253 return sigmas, averaged_sigmas

254

255 def _compute_sigma_profile(

256 self, averaged_sigmas: NDArray[np.float64], areas: NDArray[np.float64]

257 ) -> NDArray[np.float64]:

258 """Bin segment areas by charge density using linear interpolation.

259

260 Parameters

261 ----------

262 averaged_sigmas : np.ndarray

263 Averaged screening charge densities in e/Å².

264 areas : np.ndarray

265 Surface areas in Å².

266

267 Returns

268 -------

269 np.ndarray

270 Sigma profile histogram. Shape: (num_points,).

271 """

272 profile = np.zeros_like(self._grid)

273 max_index = len(self._grid) - 2 # index + 1 must be valid

274 for sigma, area in zip(averaged_sigmas, areas, strict=True):

275 index = int((sigma - self._min_sigma) / self._bin_width)

276 index = min(max(0, index), max_index)

277 weight = (self._grid[index + 1] - sigma) / self._bin_width

278 profile[index] += area * weight

279 profile[index + 1] += area * (1.0 - weight)

280 return profile

281

282 def _compute_sigma_profiles(

283 self, averaged_sigmas: NDArray[np.float64]

284 ) -> dict[str, NDArray[np.float64]]:

285 """Compute sigma profiles separated by hydrogen bonding type.

286

287 Classifies segments by H-bonding type (OH, OT, NHB) based on parent atom

288 and sigma sign, then applies a Gaussian probability weighting function.

289

290 Parameters

291 ----------

292 averaged_sigmas : np.ndarray

293 Averaged screening charge densities for all segments in e/Å².

294

295 Returns

296 -------

297 dict

298 Dictionary with keys "NHB", "OH", "OT" and values as sigma profile

299 arrays. Each profile has shape (num_points,).

300 """

301 atom_indices = self._segment_data["atom"]

302 element = atom_indices.map(self._atom_data["element"])

303 is_hb_candidate = (element == "H") == (averaged_sigmas < 0.0)

304 hb_class = atom_indices.map(self._get_hydrogen_bonding_classes())

305 mask_oh = is_hb_candidate & (hb_class == OH)

306 mask_ot = is_hb_candidate & (hb_class == OT)

307 mask_nhb = np.logical_not(mask_oh | mask_ot)

308 areas = self._segment_data["area"].values

309 profile_oh = self._compute_sigma_profile(

310 averaged_sigmas[mask_oh], areas[mask_oh]

311 )

312 profile_ot = self._compute_sigma_profile(

313 averaged_sigmas[mask_ot], areas[mask_ot]

314 )

315 profile_nhb = self._compute_sigma_profile(

316 averaged_sigmas[mask_nhb], areas[mask_nhb]

317 )

318 if self._sigma_0 is None:

319 hb_probability = np.zeros_like(self._grid)

320 else:

321 hb_probability = 1.0 - np.exp(-0.5 * (self._grid / self._sigma_0) ** 2)

322 return {

323 NHB: profile_nhb + (profile_oh + profile_ot) * (1.0 - hb_probability),

324 OH: profile_oh * hb_probability,

325 OT: profile_ot * hb_probability,

326 }

327

328 @classmethod

329 def from_text_reader(cls, text_reader: TextIO) -> "Component":

330 """Create a component from a text reader.

331

332 .. note::

333 This method creates a component with default parameters.

334

335 Parameters

336 ----------

337 text_reader : io.TextIO

338 Text reader to read the COSMO output file from.

339

340 Returns

341 -------

342 Component

343 Component object.

344

345 Examples

346 --------

347 >>> from importlib.resources import files

348 >>> from cosmolayer.cosmosac import Component

349 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

350 >>> with open(path, encoding="utf-8") as file:

351 ... component = Component.from_text_reader(file)

352 >>> component.area, component.volume

353 (97.34554..., 80.07160...)

354

355 """

356 return cls(text_reader.read())

357

358 @classmethod

359 def from_file(cls, file_path: os.PathLike[str] | Traversable) -> "Component":

360 """Create a component from a COSMO output file.

361

362 .. note::

363 This method creates a component with default parameters.

364

365 Parameters

366 ----------

367 file_path : path-like or Traversable

368 Path to the COSMO output file.

369

370 Returns

371 -------

372 Component

373 Component object.

374

375 Examples

376 --------

377 >>> from importlib.resources import files

378 >>> from cosmolayer.cosmosac import Component

379 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

380 >>> component = Component.from_file(path)

381 >>> component.area, component.volume

382 (97.34554..., 80.07160...)

383

384 """

385 if isinstance(file_path, os.PathLike):

386 with open(file_path, encoding="utf-8") as file:

387 return cls.from_text_reader(file)

388 with file_path.open("r", encoding="utf-8") as file:

389 return cls.from_text_reader(file)

390

391 @property

392 def area(self) -> float:

393 """Cavity surface area of the molecule in Å².

394

395 Sum of the areas of all segments from the COSMO calculation.

396 """

397 return self._area

398

399 @property

400 def volume(self) -> float:

401 """Cavity volume of the molecule in Å³."""

402 return self._volume

403

404 @property

405 def cosmo_format(self) -> str:

406 """COSMO file format that was parsed.

407

408 Either "TURBOMOLE" or "DMol-3".

409

410 Examples

411 --------

412 >>> from importlib.resources import files

413 >>> from cosmolayer.cosmosac import Component

414 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

415 >>> component = Component.from_file(path)

416 >>> component.cosmo_format

417 'TURBOMOLE'

418 >>> path = files("cosmolayer.data") / "NCCO.cosmo"

419 >>> component = Component.from_file(path)

420 >>> component.cosmo_format

421 'DMol-3'

422 """

423 return self._format

424

425 @property

426 def atom_data(self) -> pd.DataFrame:

427 """Atom data from the parsed COSMO file.

428

429 DataFrame columns: ``id`` (atom identifier), ``x``, ``y``, ``z`` (Cartesian

430 coordinates in Å), ``element`` (chemical symbol).

431

432 Returns

433 -------

434 pd.DataFrame

435 One row per atom.

436

437 Examples

438 --------

439 >>> from importlib.resources import files

440 >>> from cosmolayer.cosmosac import Component

441 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

442 >>> component = Component(path.read_text())

443 >>> component.atom_data

444 id x y z element

445 0 C1 -1.4... -0.2... 0.0... C

446 1 C2 -0.0... 0.0... 0.0... C

447 2 N1 0.9... -0.9... -0.0... N

448 ...

449 8 H5 1.1... 1.3... -0.4... H

450

451 """

452 return self._atom_data

453

454 @property

455 def segment_data(self) -> pd.DataFrame:

456 """Segment (surface tile) data from the COSMO calculation.

457

458 DataFrame columns: ``atom`` (parent atom index), ``x``, ``y``, ``z``

459 (segment centroid coordinates in Å), ``charge`` (e), ``area`` (Å²),

460 ``sigma`` (screening charge density in e/Å²), ``sigma_avg`` (smoothed

461 density in e/Å²).

462

463 Returns

464 -------

465 pd.DataFrame

466 One row per segment.

467

468 Examples

469 --------

470 >>> from importlib.resources import files

471 >>> from cosmolayer.cosmosac import Component

472 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

473 >>> component = Component(path.read_text())

474 >>> component.segment_data

475 atom x y ... area sigma sigma_avg

476 0 0 -0.867... -1.196... ... 0.206... 0.010... 0.007...

477 1 0 -1.504... -1.502... ... 0.218... 0.007... 0.005...

478 ...

479 470 8 2.133... 1.152... ... 0.145... -0.012... -0.009...

480 <BLANKLINE>

481 [471 rows x 8 columns]

482

483 """

484 return self._segment_data

485

486 @property

487 def bonds(self) -> list[tuple[int, int]]:

488 """Bonds between atoms, inferred from interatomic distances.

489

490 Returns

491 -------

492 list[tuple[int, int]]

493 Pairs of atom indices (i, j) for each bond.

494

495 Examples

496 --------

497 >>> from importlib.resources import files

498 >>> from cosmolayer.cosmosac import Component

499 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

500 >>> component = Component(path.read_text())

501 >>> component.bonds

502 [(0, 1), (0, 4), (0, 5), ... (2, 7), (3, 8)]

503 """

504 return self._bonds

505

506 @property

507 def sigma_grid(self) -> NDArray[np.float64]:

508 """Get the screening charge density grid in e/Å².

509

510 Returns

511 -------

512 np.ndarray

513 Charge density vector in e/Å².

514

515 Examples

516 --------

517 >>> from importlib.resources import files

518 >>> from cosmolayer.cosmosac import Component

519 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"

520 >>> component = Component(path.read_text())

521 >>> component.sigma_grid

522 array([-0.025, -0.024, -0.023, ... 0.023, 0.024, 0.025])

523 """

524 return self._grid

525

526 @property

527 def merge_profiles(self) -> bool:

528 """Whether segment groups (NHB, OH, OT) are merged for :attr:`sigma_profile`

529 and :attr:`probabilities`.

530

531 Returns

532 -------

533 bool

534 """

535 return self._merge_profiles

536

537 @property

538 def sigma_profile(self) -> NDArray[np.float64]:

539 """Surface area distribution over screening charge density (sigma), in Å².

540

541 Shape and layout depend on :attr:`merge_profiles`. If True, returns a single

542 merged profile (sum over NHB, OH, OT), shape ``(num_points,)``. If False,

543 returns stacked segment profiles in SEGMENT_GROUPS order (NHB, OH, OT),

544 shape ``(3, num_points)``; ``sigma_profile[0]`` is NHB, ``[1]`` is OH,

545 ``[2]`` is OT.

546

547 Returns

548 -------

549 np.ndarray

550 Sigma profile(s). Units: Å².

551 """

552 if self._merge_profiles:

553 total_profile: NDArray[np.float64] = np.sum(

554 list(self._sigma_profiles.values()), axis=0

555 )

556 return total_profile

557 return np.stack([self._sigma_profiles[seg] for seg in SEGMENT_GROUPS], axis=0)

558

559 @property

560 def probabilities(self) -> NDArray[np.float64]:

561 """Normalized segment-type probability distribution (sigma profile / area).

562

563 Segment types are defined by hydrogen bonding class (NHB, OH, OT) and

564 averaged charge density. Shape is ``(num_points,)`` if :attr:`merge_profiles`

565 is True, otherwise ``(3*num_points,)``.

566

567 Returns

568 -------

569 np.ndarray

570 Probabilities summing to 1.0.

571

572 Examples

573 --------

574 >>> import numpy as np

575 >>> from importlib.resources import files

576 >>> from cosmolayer.cosmosac import Component

577 >>> cosmo_string = (files("cosmolayer.data") / "C=C(N)O.cosmo").read_text()

578 >>> component = Component(cosmo_string, merge_profiles=True)

579 >>> probabilities = component.probabilities

580 >>> probabilities.shape

581 (51,)

582 >>> bool(np.all(probabilities <= 1))

583 True

584 >>> bool(np.isclose(probabilities.sum(), 1.0))

585 True

586 >>> component = Component(cosmo_string, merge_profiles=False)

587 >>> probabilities_full = component.probabilities

588 >>> probabilities_full.shape

589 (153,)

590 >>> bool(np.isclose(probabilities_full.sum(), 1.0))

591 True

592 """

593 profiles = [self._sigma_profiles[segtype] for segtype in SEGMENT_GROUPS]

594 probabilities: NDArray[np.float64] = (

595 np.sum(profiles, axis=0)

596 if self._merge_profiles

597 else np.concatenate(profiles)

598 ) / self._area

599 return probabilities