# This file is part of tad-dftd3.
# SPDX-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Reference model
===============
This module defines the reference systems for the D3 model to compute the
C6 dispersion coefficients.
"""
import os.path as op
from typing import Optional
import torch
from .typing import Any, NoReturn, Tensor, get_default_device, get_default_dtype
__all__ = ["Reference"]
def _load_cn(
dtype: torch.dtype = torch.double, device: Optional[torch.device] = None
) -> Tensor:
"""
Load reference coordination numbers.
Parameters
----------
dtype : torch.dtype, optional
Floating point precision for tensor. Defaults to `torch.double`.
device : Optional[torch.device], optional
Device of tensor. Defaults to None.
Returns
-------
Tensor
Reference coordination numbers.
"""
return torch.tensor(
[
[-1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # None
[+0.9118, +0.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # H
[+0.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # He
[+0.0000, +0.9865, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Li
[+0.0000, +0.9808, +1.9697, -1.0000, -1.0000, -1.0000, -1.0000], # Be
[+0.0000, +0.9706, +1.9441, +2.9128, +4.5856, -1.0000, -1.0000], # B
[+0.0000, +0.9868, +1.9985, +2.9987, +3.9844, -1.0000, -1.0000], # C
[+0.0000, +0.9944, +2.0143, +2.9903, -1.0000, -1.0000, -1.0000], # N
[+0.0000, +0.9925, +1.9887, -1.0000, -1.0000, -1.0000, -1.0000], # O
[+0.0000, +0.9982, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # F
[+0.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Ne
[+0.0000, +0.9684, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Na
[+0.0000, +0.9628, +1.9496, -1.0000, -1.0000, -1.0000, -1.0000], # Mg
[+0.0000, +0.9648, +1.9311, +2.9146, -1.0000, -1.0000, -1.0000], # Al
[+0.0000, +0.9507, +1.9435, +2.9407, +3.8677, -1.0000, -1.0000], # Si
[+0.0000, +0.9947, +2.0102, +2.9859, -1.0000, -1.0000, -1.0000], # P
[+0.0000, +0.9948, +1.9903, -1.0000, -1.0000, -1.0000, -1.0000], # S
[+0.0000, +0.9972, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Cl
[+0.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Ar
[+0.0000, +0.9767, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # K
[+0.0000, +0.9831, +1.9349, -1.0000, -1.0000, -1.0000, -1.0000], # Ca
[+0.0000, +1.8627, +2.8999, -1.0000, -1.0000, -1.0000, -1.0000], # Sc
[+0.0000, +1.8299, +3.8675, -1.0000, -1.0000, -1.0000, -1.0000], # Ti
[+0.0000, +1.9138, +2.9110, -1.0000, -1.0000, -1.0000, -1.0000], # V
[+0.0000, +1.8269, 10.6191, -1.0000, -1.0000, -1.0000, -1.0000], # Cr
[+0.0000, +1.6406, +9.8849, -1.0000, -1.0000, -1.0000, -1.0000], # Mn
[+0.0000, +1.6483, +9.1376, -1.0000, -1.0000, -1.0000, -1.0000], # Fe
[+0.0000, +1.7149, +2.9263, +7.7785, -1.0000, -1.0000, -1.0000], # Co
[+0.0000, +1.7937, +6.5458, +6.2918, -1.0000, -1.0000, -1.0000], # Ni
[+0.0000, +0.9576, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Cu
[+0.0000, +1.9419, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Zn
[+0.0000, +0.9601, +1.9315, +2.9233, -1.0000, -1.0000, -1.0000], # Ga
[+0.0000, +0.9434, +1.9447, +2.9186, +3.8972, -1.0000, -1.0000], # Ge
[+0.0000, +0.9889, +1.9793, +2.9709, -1.0000, -1.0000, -1.0000], # As
[+0.0000, +0.9901, +1.9812, -1.0000, -1.0000, -1.0000, -1.0000], # Se
[+0.0000, +0.9974, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Br
[+0.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Kr
[+0.0000, +0.9738, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Rb
[+0.0000, +0.9801, +1.9143, -1.0000, -1.0000, -1.0000, -1.0000], # Sr
[+0.0000, +1.9153, +2.8903, -1.0000, -1.0000, -1.0000, -1.0000], # Y
[+0.0000, +1.9355, +3.9106, -1.0000, -1.0000, -1.0000, -1.0000], # Zr
[+0.0000, +1.9545, +2.9225, -1.0000, -1.0000, -1.0000, -1.0000], # Nb
[+0.0000, +1.9420, 11.0556, -1.0000, -1.0000, -1.0000, -1.0000], # Mo
[+0.0000, +1.6682, +9.5402, -1.0000, -1.0000, -1.0000, -1.0000], # Tc
[+0.0000, +1.8584, +8.8895, -1.0000, -1.0000, -1.0000, -1.0000], # Ru
[+0.0000, +1.9003, +2.9696, -1.0000, -1.0000, -1.0000, -1.0000], # Rh
[+0.0000, +1.8630, +5.7095, -1.0000, -1.0000, -1.0000, -1.0000], # Pd
[+0.0000, +0.9679, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Ag
[+0.0000, +1.9539, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Cd
[+0.0000, +0.9633, +1.9378, +2.9353, -1.0000, -1.0000, -1.0000], # In
[+0.0000, +0.9514, +1.9505, +2.9259, +3.9123, -1.0000, -1.0000], # Sn
[+0.0000, +0.9749, +1.9523, +2.9315, -1.0000, -1.0000, -1.0000], # Sb
[+0.0000, +0.9811, +1.9639, -1.0000, -1.0000, -1.0000, -1.0000], # Te
[+0.0000, +0.9968, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # I
[+0.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Xe
[+0.0000, +0.9909, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Cs
[+0.0000, +0.9797, +1.8467, -1.0000, -1.0000, -1.0000, -1.0000], # Ba
[+0.0000, +1.9373, +2.9175, -1.0000, -1.0000, -1.0000, -1.0000], # La
[+2.7991, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Ce
[+0.0000, +2.9425, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Pr
[+0.0000, +2.9455, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Nd
[+0.0000, +2.9413, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Pm
[+0.0000, +2.9300, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Sm
[+0.0000, +1.8286, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Eu
[+0.0000, +2.8732, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Gd
[+0.0000, +2.9086, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Tb
[+0.0000, +2.8965, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Dy
[+0.0000, +2.9242, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Ho
[+0.0000, +2.9282, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Er
[+0.0000, +2.9246, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Tm
[+0.0000, +2.8482, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Yb
[+0.0000, +2.9219, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Lu
[+0.0000, +1.9254, +3.8840, -1.0000, -1.0000, -1.0000, -1.0000], # Hf
[+0.0000, +1.9459, +2.8988, -1.0000, -1.0000, -1.0000, -1.0000], # Ta
[+0.0000, +1.9292, 10.9153, -1.0000, -1.0000, -1.0000, -1.0000], # W
[+0.0000, +1.8104, +9.8054, -1.0000, -1.0000, -1.0000, -1.0000], # Re
[+0.0000, +1.8858, +9.1527, -1.0000, -1.0000, -1.0000, -1.0000], # Os
[+0.0000, +1.8648, +2.9424, -1.0000, -1.0000, -1.0000, -1.0000], # Ir
[+0.0000, +1.9188, +6.6669, -1.0000, -1.0000, -1.0000, -1.0000], # Pt
[+0.0000, +0.9846, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Au
[+0.0000, +1.9896, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Hg
[+0.0000, +0.9267, +1.9302, +2.9420, -1.0000, -1.0000, -1.0000], # Tl
[+0.0000, +0.9383, +1.9356, +2.9081, +3.9098, -1.0000, -1.0000], # Pb
[+0.0000, +0.9820, +1.9655, +2.9500, -1.0000, -1.0000, -1.0000], # Bi
[+0.0000, +0.9815, +1.9639, -1.0000, -1.0000, -1.0000, -1.0000], # Po
[+0.0000, +0.9954, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # At
[+0.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Rn
[+0.0000, +0.9705, -1.0000, -1.0000, -1.0000, -1.0000, -1.0000], # Fr
[+0.0000, +0.9661, +1.9251, -1.0000, -1.0000, -1.0000, -1.0000], # Ra
[+0.0000, +0.9802, +1.9445, +2.9070, +3.8174, +4.6723, +5.5599], # Ac
[+0.0000, +0.9847, +1.9560, +2.9302, +3.8997, -1.0000, -1.0000], # Th
[+0.0000, +0.9647, +1.9079, +2.9037, +3.8711, +4.9094, +4.5318], # Pa
[+0.0000, +0.9766, +2.8888, +3.9129, +4.1181, +5.9187, -1.0000], # U
[+0.0000, +0.9838, +1.9499, +2.9159, +3.9358, +4.9069, +5.9005], # Np
[+0.0000, +0.9537, +1.9439, +2.9323, +3.9441, +4.9192, +5.8888], # Pu
[+0.0000, +0.9163, +1.8563, +2.8823, +4.8005, +5.7794, -1.0000], # Am
[+0.0000, +0.9762, +1.9288, +2.8929, +3.8167, +4.7478, +5.6866], # Cm
[+0.0000, +0.9705, +1.9511, +2.9262, +3.9342, -1.0000, -1.0000], # Bk
[+0.0000, +0.9581, +1.9123, +2.9327, +3.9105, +5.8285, -1.0000], # Cf
[+0.0000, +0.9346, +1.8816, +2.9075, +3.8705, +4.8131, +5.7244], # Es
[+0.0000, +0.9500, +1.9165, +2.9377, +3.8956, +4.8540, +5.8160], # Fm
[+0.0000, +0.9710, +1.9564, +2.9515, +3.9353, -1.0000, -1.0000], # Md
[+0.0000, +0.9722, +1.9605, +2.9452, +3.9296, +4.2582, +4.5511], # No
[+0.0000, +0.9569, +1.9215, +2.8958, +3.7644, +4.6808, +5.5939], # Lr
],
device=device,
dtype=dtype,
)
def _load_c6(
dtype: torch.dtype = torch.double, device: Optional[torch.device] = None
) -> Tensor:
"""
Load reference C6 coefficients from file.
Parameters
----------
dtype : torch.dtype, optional
Floating point precision for tensor. Defaults to `torch.double`.
device : Optional[torch.device], optional
Device of tensor. Defaults to None.
Returns
-------
Tensor
Reference C6 coefficients.
"""
path = op.join(op.dirname(__file__), "reference-c6.pt")
return torch.load(path).type(dtype).to(device)
[docs]
class Reference:
"""
Reference systems for the D3 dispersion model
"""
c6: Tensor
"""C6 coefficients for all pairs of reference systems"""
cn: Tensor
"""Coordination numbers for all reference systems"""
__slots__ = [
"c6",
"cn",
"__dtype",
"__device",
]
def __init__(
self,
cn: Optional[Tensor] = None,
c6: Optional[Tensor] = None,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
):
if cn is None:
cn = _load_cn(
dtype=dtype if dtype is not None else get_default_dtype(),
device=device if device is not None else get_default_device(),
)
self.cn = cn
if c6 is None:
c6 = _load_c6(
dtype=dtype if dtype is not None else get_default_dtype(),
device=device if device is not None else get_default_device(),
)
self.c6 = c6
self.__dtype = self.c6.dtype
self.__device = self.c6.device
if any(tensor.device != self.device for tensor in (self.cn, self.c6)):
raise RuntimeError("All tensors must be on the same device!")
if any(tensor.dtype != self.dtype for tensor in (self.cn, self.c6)):
raise RuntimeError("All tensors must have the same dtype!")
if any(
(
self.c6.shape[-2] != self.c6.shape[-1],
self.c6.shape[-1] != self.cn.shape[-1],
self.c6.shape[-4] != self.c6.shape[-3],
self.c6.shape[-3] != self.cn.shape[-2],
)
):
raise RuntimeError("`c6` & `cn` size mismatch found")
@property
def device(self) -> torch.device:
"""The device on which the `Reference` object resides."""
return self.__device
@device.setter
def device(self, *_: Any) -> NoReturn:
"""
Instruct users to use the ".to" method if wanting to change device.
"""
raise AttributeError("Move object to device using the `.to` method")
@property
def dtype(self) -> torch.dtype:
"""Floating point dtype used by reference object."""
return self.__dtype
[docs]
def to(
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
) -> "Reference":
"""
Returns a copy of the `Reference` instance on the specified device.
This method creates and returns a new copy of the `Reference` instance
on the specified device "``device``".
Parameters
----------
device : torch.device, optional
Device to which all associated tensors should be moved.
dtype : torch.dtype, optional
Floating point type of the tensors.
Returns
-------
Reference
A copy of the `Reference` instance placed on the specified device.
Notes
-----
If the `Reference` instance is already on the desired device `self`
will be returned.
"""
if self.__device == device:
if dtype is not None:
return self.type(dtype)
return self
return self.__class__(
self.cn.to(device=device, dtype=dtype),
self.c6.to(device=device, dtype=dtype),
)
[docs]
def type(self, dtype: torch.dtype) -> "Reference":
"""
Returns a copy of the `Reference` instance with specified floating
point type. This method creates and returns a new copy of the
`Reference` instance with the specified dtype.
Parameters
----------
dtype : torch.dtype
Floating point type of the tensors.
Returns
-------
Reference
A copy of the `Reference` instance with the specified dtype.
Notes
-----
If the `Reference` instance has already the desired dtype `self` will
be returned.
"""
if self.__dtype == dtype:
return self
return self.__class__(
self.cn.type(dtype),
self.c6.type(dtype),
)
def __str__(self) -> str:
"""Creates a string representation of the Reference object."""
return (
f"{self.__class__.__name__}(n_element={self.cn.shape[-2]}, "
f"n_reference={self.cn.shape[-1]}, dtype={self.__dtype}, "
f"device={self.__device})"
)
def __repr__(self) -> str:
"""Creates a string representation of the Reference object."""
return str(self)
