GraphAtoms 0.0.1__tar.gz → 0.0.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- graphatoms-0.0.2/GraphAtoms/__init__.py +8 -0
- graphatoms-0.0.2/GraphAtoms/common/__init__.py +4 -0
- graphatoms-0.0.2/GraphAtoms/common/abc.py +285 -0
- graphatoms-0.0.2/GraphAtoms/common/error.py +16 -0
- graphatoms-0.0.2/GraphAtoms/common/geometry.py +376 -0
- graphatoms-0.0.2/GraphAtoms/common/ndarray.py +3 -0
- graphatoms-0.0.2/GraphAtoms/common/ndarray.pyi +5 -0
- graphatoms-0.0.2/GraphAtoms/common/rotation.py +101 -0
- graphatoms-0.0.2/GraphAtoms/common/string.py +33 -0
- graphatoms-0.0.2/GraphAtoms/containner/__init__.py +20 -0
- graphatoms-0.0.2/GraphAtoms/containner/_atomic.py +366 -0
- graphatoms-0.0.2/GraphAtoms/containner/_bonded.py +90 -0
- graphatoms-0.0.2/GraphAtoms/containner/_graph.py +704 -0
- graphatoms-0.0.2/GraphAtoms/containner/_mConnComp.py +45 -0
- graphatoms-0.0.2/GraphAtoms/containner/_mDataPyG.py +87 -0
- graphatoms-0.0.2/GraphAtoms/containner/_mFreeE.py +117 -0
- graphatoms-0.0.2/GraphAtoms/containner/_mIGraph.py +269 -0
- graphatoms-0.0.2/GraphAtoms/containner/_mNetworKit.py +12 -0
- graphatoms-0.0.2/GraphAtoms/containner/_mNetworkX.py +12 -0
- graphatoms-0.0.2/GraphAtoms/containner/_mRDKit.py +337 -0
- graphatoms-0.0.2/GraphAtoms/containner/_mRustworkX.py +12 -0
- graphatoms-0.0.2/GraphAtoms/containner/_sysCluster.py +121 -0
- graphatoms-0.0.2/GraphAtoms/containner/_sysGas.py +81 -0
- graphatoms-0.0.2/GraphAtoms/containner/_system.py +15 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/GraphAtoms.egg-info/PKG-INFO +1 -2
- graphatoms-0.0.2/GraphAtoms.egg-info/SOURCES.txt +36 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/GraphAtoms.egg-info/requires.txt +0 -1
- {graphatoms-0.0.1 → graphatoms-0.0.2}/PKG-INFO +1 -2
- {graphatoms-0.0.1 → graphatoms-0.0.2}/pyproject.toml +1 -7
- graphatoms-0.0.1/GraphAtoms.egg-info/SOURCES.txt +0 -12
- {graphatoms-0.0.1 → graphatoms-0.0.2}/GraphAtoms.egg-info/dependency_links.txt +0 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/GraphAtoms.egg-info/top_level.txt +0 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/LICENSE +0 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/README.md +0 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/setup.cfg +0 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/tests/test_abc.py +0 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/tests/test_container.py +0 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/tests/test_mRDKit.py +0 -0
- {graphatoms-0.0.1 → graphatoms-0.0.2}/tests/test_utils.py +0 -0
|
@@ -0,0 +1,285 @@
|
|
|
1
|
+
# ruff: noqa: D100 D102
|
|
2
|
+
from collections.abc import Sequence
|
|
3
|
+
from typing import override
|
|
4
|
+
|
|
5
|
+
import numpy as np
|
|
6
|
+
import pydantic
|
|
7
|
+
from joblib import dump, load
|
|
8
|
+
from numpy.typing import ArrayLike
|
|
9
|
+
from typing_extensions import Self
|
|
10
|
+
|
|
11
|
+
__all__ = [
|
|
12
|
+
"BaseModel",
|
|
13
|
+
"NpzPklBaseModel",
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class BaseMixin:
|
|
18
|
+
@staticmethod
|
|
19
|
+
def get_mask_or_index(k: ArrayLike, n: int) -> np.ndarray:
|
|
20
|
+
if np.isscalar(k):
|
|
21
|
+
if isinstance(k, bool):
|
|
22
|
+
k = [k] * n
|
|
23
|
+
elif isinstance(k, int):
|
|
24
|
+
k = [k]
|
|
25
|
+
else:
|
|
26
|
+
raise TypeError(f"Unsupported type({type(k)}): {k}.")
|
|
27
|
+
subset = np.asarray(k)
|
|
28
|
+
if subset.dtype == bool:
|
|
29
|
+
if subset.size != n:
|
|
30
|
+
raise KeyError(
|
|
31
|
+
f"Except {n} boolean array, but {subset.size} got."
|
|
32
|
+
)
|
|
33
|
+
else:
|
|
34
|
+
subset = np.unique(subset.astype(int))
|
|
35
|
+
if not 0 <= max(subset) < n:
|
|
36
|
+
raise KeyError(
|
|
37
|
+
f"Except 0-{n - 1} integer array, but the "
|
|
38
|
+
f"max({max(subset)}),min({min(subset)}) got."
|
|
39
|
+
)
|
|
40
|
+
return subset.flatten()
|
|
41
|
+
|
|
42
|
+
@classmethod
|
|
43
|
+
def get_index(cls, k: ArrayLike, n: int) -> np.ndarray:
|
|
44
|
+
result = cls.get_mask_or_index(k=k, n=n)
|
|
45
|
+
if result.dtype == bool:
|
|
46
|
+
result = np.arange(n)[result]
|
|
47
|
+
return result.astype(int, copy=False)
|
|
48
|
+
|
|
49
|
+
@classmethod
|
|
50
|
+
def get_mask(cls, k: ArrayLike, n: int) -> np.ndarray:
|
|
51
|
+
result = cls.get_mask_or_index(k=k, n=n)
|
|
52
|
+
if result.dtype != bool:
|
|
53
|
+
result = np.isin(np.arange(n), result)
|
|
54
|
+
return result.astype(bool, copy=False)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
class BaseModel(pydantic.BaseModel):
|
|
58
|
+
"""A base class for creating Pydantic models.
|
|
59
|
+
|
|
60
|
+
This class support many file formats:
|
|
61
|
+
json by `pydantic`
|
|
62
|
+
"""
|
|
63
|
+
|
|
64
|
+
model_config = pydantic.ConfigDict(frozen=True)
|
|
65
|
+
|
|
66
|
+
@override
|
|
67
|
+
def __str__(self) -> str:
|
|
68
|
+
return f"{self.__class__.__name__}({self._string()})"
|
|
69
|
+
|
|
70
|
+
def _string(self) -> str:
|
|
71
|
+
"""The string representation of this class.
|
|
72
|
+
|
|
73
|
+
The expression like `str(obj)` will output
|
|
74
|
+
`f"{self.__class__.__name__}({this_value})"`.
|
|
75
|
+
"""
|
|
76
|
+
return "object"
|
|
77
|
+
|
|
78
|
+
@staticmethod
|
|
79
|
+
@pydantic.validate_call
|
|
80
|
+
def __get_format(fname: pydantic.FilePath | pydantic.NewPath) -> str:
|
|
81
|
+
format: str = fname.name.split(".")[-1].lower()
|
|
82
|
+
if format == "yml":
|
|
83
|
+
format = "yaml"
|
|
84
|
+
elif format == "pkl":
|
|
85
|
+
format = "pickle"
|
|
86
|
+
return format
|
|
87
|
+
|
|
88
|
+
def write(self, fname: pydantic.FilePath, **kwargs) -> pydantic.FilePath:
|
|
89
|
+
f = getattr(self, f"write_{self.__get_format(fname)}")
|
|
90
|
+
return f(fname, **kwargs)
|
|
91
|
+
|
|
92
|
+
@pydantic.validate_call
|
|
93
|
+
def write_json(
|
|
94
|
+
self,
|
|
95
|
+
filename: pydantic.FilePath | pydantic.NewPath,
|
|
96
|
+
indent: int | None = 4,
|
|
97
|
+
**kwargs,
|
|
98
|
+
) -> pydantic.FilePath:
|
|
99
|
+
filename.write_text(
|
|
100
|
+
self.model_dump_json(indent=indent, **kwargs),
|
|
101
|
+
encoding="utf-8",
|
|
102
|
+
)
|
|
103
|
+
return filename
|
|
104
|
+
|
|
105
|
+
@classmethod
|
|
106
|
+
@pydantic.validate_call
|
|
107
|
+
def read_json(cls, filename: pydantic.FilePath, **kwargs) -> Self:
|
|
108
|
+
return cls.model_validate_json(filename.read_bytes(), **kwargs)
|
|
109
|
+
|
|
110
|
+
@classmethod
|
|
111
|
+
def read(cls, fname: pydantic.FilePath, **kwargs) -> Self:
|
|
112
|
+
f = getattr(cls, f"read_{cls.__get_format(fname)}")
|
|
113
|
+
return f(fname, **kwargs)
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
class NpzPklBaseModel(BaseModel):
|
|
117
|
+
"""A extended base class for creating Pydantic models.
|
|
118
|
+
|
|
119
|
+
This class support many file formats:
|
|
120
|
+
json by `pydantic`
|
|
121
|
+
pickle by `joblib`
|
|
122
|
+
npz by `numpy`
|
|
123
|
+
|
|
124
|
+
Note: only numpy ndarray & numpy-compatible scalar value supported.
|
|
125
|
+
"""
|
|
126
|
+
|
|
127
|
+
@classmethod
|
|
128
|
+
def _convert(cls) -> dict[str, tuple[tuple, str]]:
|
|
129
|
+
"""Override if needed to specify the dtype and shape of attributes.
|
|
130
|
+
|
|
131
|
+
Note: must be implemented by children class.
|
|
132
|
+
|
|
133
|
+
Example:
|
|
134
|
+
return dict(a=((None, 3), "uint8"))
|
|
135
|
+
This `None` indicates it can take any shape along the first axis.
|
|
136
|
+
"""
|
|
137
|
+
return {}
|
|
138
|
+
|
|
139
|
+
@classmethod
|
|
140
|
+
def _npz_available(cls) -> bool:
|
|
141
|
+
return True
|
|
142
|
+
|
|
143
|
+
@staticmethod
|
|
144
|
+
def __validate_ndarray_and_convert( # noqa: D103
|
|
145
|
+
data: np.ndarray | Sequence,
|
|
146
|
+
shape: Sequence[int | None],
|
|
147
|
+
dtype: str,
|
|
148
|
+
) -> np.ndarray:
|
|
149
|
+
assert shape.count(None) <= 1, (type(shape), shape)
|
|
150
|
+
shape = tuple(int(i) if i is not None else -1 for i in shape)
|
|
151
|
+
if not isinstance(data, np.ndarray):
|
|
152
|
+
data = np.asarray(data)
|
|
153
|
+
return np.asarray(data, dtype).reshape(shape)
|
|
154
|
+
|
|
155
|
+
@pydantic.model_validator(mode="before")
|
|
156
|
+
@classmethod
|
|
157
|
+
def __convert(cls, data):
|
|
158
|
+
if isinstance(data, dict):
|
|
159
|
+
for key, (shape, dtype) in cls._convert().items():
|
|
160
|
+
if key in data and data[key] is not None:
|
|
161
|
+
kw = dict(data=data[key], shape=shape, dtype=dtype)
|
|
162
|
+
data[key] = cls.__validate_ndarray_and_convert(**kw) # type: ignore
|
|
163
|
+
return data
|
|
164
|
+
|
|
165
|
+
@pydantic.model_validator(mode="after")
|
|
166
|
+
def __check_shape(self) -> Self:
|
|
167
|
+
for k in self.__pydantic_fields__:
|
|
168
|
+
v = getattr(self, k, None)
|
|
169
|
+
if v is None or np.isscalar(v) or isinstance(v, np.ndarray):
|
|
170
|
+
pass
|
|
171
|
+
else:
|
|
172
|
+
raise ValueError(f"Invalid value: {v}")
|
|
173
|
+
return self
|
|
174
|
+
|
|
175
|
+
@override
|
|
176
|
+
def __eq__(self, other) -> bool:
|
|
177
|
+
if not isinstance(other, self.__class__):
|
|
178
|
+
return False
|
|
179
|
+
for k in self.__pydantic_fields__:
|
|
180
|
+
v0 = getattr(self, k, None)
|
|
181
|
+
v1 = getattr(other, k, None)
|
|
182
|
+
if v0 is None or np.isscalar(v0):
|
|
183
|
+
if v0 is None or np.isscalar(v0):
|
|
184
|
+
if v0 != v1:
|
|
185
|
+
return False
|
|
186
|
+
else:
|
|
187
|
+
return False
|
|
188
|
+
elif isinstance(v0, np.ndarray):
|
|
189
|
+
if isinstance(v1, np.ndarray):
|
|
190
|
+
if v0.shape != v1.shape:
|
|
191
|
+
return False
|
|
192
|
+
elif np.any(v0 != v1):
|
|
193
|
+
return False
|
|
194
|
+
else:
|
|
195
|
+
return False
|
|
196
|
+
else:
|
|
197
|
+
raise NotImplementedError(f"Field({k}): {type(v0)}")
|
|
198
|
+
return True
|
|
199
|
+
|
|
200
|
+
@pydantic.validate_call
|
|
201
|
+
def write_pickle(
|
|
202
|
+
self,
|
|
203
|
+
filename: pydantic.FilePath | pydantic.NewPath,
|
|
204
|
+
compress: bool | int | tuple[str, int] = 3,
|
|
205
|
+
**kwargs,
|
|
206
|
+
) -> pydantic.FilePath:
|
|
207
|
+
"""Persist the dictionary of this object into one file.
|
|
208
|
+
|
|
209
|
+
Read more in the reference:
|
|
210
|
+
https://joblib.readthedocs.io/en/latest/generated/joblib.dump.html
|
|
211
|
+
|
|
212
|
+
Args:
|
|
213
|
+
filename (str | Path | TextIOWrapper): The file object or path
|
|
214
|
+
of the file in which it is to be stored. The compression
|
|
215
|
+
method corresponding to one of the supported filename
|
|
216
|
+
extensions ('.z', '.gz', '.bz2', '.xz' or '.lzma')
|
|
217
|
+
will be used automatically.
|
|
218
|
+
compress (bool | int | tuple[str, int], optional):
|
|
219
|
+
Optional compression level for the data.
|
|
220
|
+
0 or False is no compression.
|
|
221
|
+
Higher value means more compression, but also
|
|
222
|
+
slower read and write times.
|
|
223
|
+
Using a value of 3 is often a good compromise.
|
|
224
|
+
Defaults to 3.
|
|
225
|
+
exclude_defaults (bool): ...
|
|
226
|
+
exclude_none (bool): ...
|
|
227
|
+
**kwargs: ...
|
|
228
|
+
|
|
229
|
+
Note:
|
|
230
|
+
If compress is True, the compression level used is 3.
|
|
231
|
+
If compress is a 2-tuple, the first element must correspond to
|
|
232
|
+
a string between supported compressors (e.g 'zlib', 'gzip',
|
|
233
|
+
'bz2', 'lzma' 'xz'), the second element must be an integer
|
|
234
|
+
from 0 to 9, corresponding to the compression level.
|
|
235
|
+
"""
|
|
236
|
+
kwargs["exclude_none"] = True
|
|
237
|
+
kwargs["exclude_defaults"] = False
|
|
238
|
+
dump(
|
|
239
|
+
self.model_dump(mode="python", **kwargs),
|
|
240
|
+
filename,
|
|
241
|
+
compress=compress, # type: ignore
|
|
242
|
+
)
|
|
243
|
+
return filename
|
|
244
|
+
|
|
245
|
+
@pydantic.validate_call
|
|
246
|
+
def write_npz(
|
|
247
|
+
self,
|
|
248
|
+
filename: pydantic.FilePath | pydantic.NewPath,
|
|
249
|
+
compress: bool = True,
|
|
250
|
+
**kwargs,
|
|
251
|
+
) -> pydantic.FilePath:
|
|
252
|
+
if not self._npz_available():
|
|
253
|
+
raise RuntimeError(
|
|
254
|
+
"Subclass must have class method named"
|
|
255
|
+
" `_npz_available()` to be True."
|
|
256
|
+
)
|
|
257
|
+
kwargs["exclude_none"] = True
|
|
258
|
+
kwargs["exclude_defaults"] = False
|
|
259
|
+
f_savez = np.savez_compressed if compress else np.savez
|
|
260
|
+
data = self.model_dump(mode="python", **kwargs)
|
|
261
|
+
dct: dict[str, ArrayLike] = {}
|
|
262
|
+
for k, v in data.items():
|
|
263
|
+
if not isinstance(v, dict):
|
|
264
|
+
dct[k] = v
|
|
265
|
+
f_savez(
|
|
266
|
+
filename,
|
|
267
|
+
allow_pickle=False,
|
|
268
|
+
**self.model_dump(mode="python", **kwargs),
|
|
269
|
+
)
|
|
270
|
+
return filename
|
|
271
|
+
|
|
272
|
+
@classmethod
|
|
273
|
+
@pydantic.validate_call
|
|
274
|
+
def read_npz(cls, filename: pydantic.FilePath, **kwargs) -> Self:
|
|
275
|
+
if not cls._npz_available():
|
|
276
|
+
raise RuntimeError(
|
|
277
|
+
"Subclass must have class method named"
|
|
278
|
+
" `_npz_available()` to be True."
|
|
279
|
+
)
|
|
280
|
+
return cls.model_validate(np.load(filename), **kwargs)
|
|
281
|
+
|
|
282
|
+
@classmethod
|
|
283
|
+
@pydantic.validate_call
|
|
284
|
+
def read_pickle(cls, filename: pydantic.FilePath, **kwargs) -> Self:
|
|
285
|
+
return cls.model_validate(load(filename), **kwargs)
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
# ruff: noqa: D100 D101 D102 D103
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class NotSupportPeriodicSystem(KeyError):
|
|
5
|
+
"""The periodic system is not supported."""
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class NotSupportNonOrthorhombicLattice(KeyError):
|
|
9
|
+
"""The non-orthorhombic lattice is not supported."""
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def not_support_type(data, required_type: str) -> TypeError:
|
|
13
|
+
return TypeError(
|
|
14
|
+
f"Only support type({required_type:s}). But"
|
|
15
|
+
f" the type of input is {type(data)} now."
|
|
16
|
+
)
|
|
@@ -0,0 +1,376 @@
|
|
|
1
|
+
"""The calculation of distance."""
|
|
2
|
+
|
|
3
|
+
import itertools
|
|
4
|
+
from typing import Literal
|
|
5
|
+
|
|
6
|
+
import numpy as np
|
|
7
|
+
import sparse
|
|
8
|
+
from ase.cell import Cell
|
|
9
|
+
from ase.data import covalent_radii as COV_R
|
|
10
|
+
from ase.geometry import find_mic, minkowski_reduce, wrap_positions
|
|
11
|
+
from scipy import sparse as sp
|
|
12
|
+
from scipy.linalg import norm, pinv
|
|
13
|
+
from sklearn.neighbors import NearestNeighbors
|
|
14
|
+
|
|
15
|
+
from GraphAtoms.common.rotation import kabsch, rotate
|
|
16
|
+
|
|
17
|
+
rigid_transform_3D = kabsch
|
|
18
|
+
__all__ = [
|
|
19
|
+
"kabsch",
|
|
20
|
+
"rotate",
|
|
21
|
+
"rigid_transform_3D",
|
|
22
|
+
"get_distance_sparse_matrix",
|
|
23
|
+
"get_adjacency_sparse_matrix",
|
|
24
|
+
"get_is_inner_array",
|
|
25
|
+
"rigid_transform_3D",
|
|
26
|
+
"distance_factory",
|
|
27
|
+
]
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
class _DistanceFactory:
|
|
31
|
+
"""The factory of distance-related calculation."""
|
|
32
|
+
|
|
33
|
+
@staticmethod
|
|
34
|
+
def __get_nn(X) -> NearestNeighbors:
|
|
35
|
+
nn = NearestNeighbors()
|
|
36
|
+
nn.fit(X=X, y=None)
|
|
37
|
+
return nn
|
|
38
|
+
|
|
39
|
+
@staticmethod
|
|
40
|
+
def __get_sparse_matrix(
|
|
41
|
+
nn: NearestNeighbors,
|
|
42
|
+
points: np.ndarray,
|
|
43
|
+
radius: float,
|
|
44
|
+
return_distance: bool = False,
|
|
45
|
+
) -> sp.csr_matrix:
|
|
46
|
+
"""Get sparse matrix by distance.
|
|
47
|
+
|
|
48
|
+
Args:
|
|
49
|
+
nn (NearestNeighbors): The sklearn NearestNeighbors object.
|
|
50
|
+
points (np.ndarray): The query point or points.
|
|
51
|
+
radius (float): Radius of neighborhoods.
|
|
52
|
+
return_distance (bool, optional): Defaults to False.
|
|
53
|
+
|
|
54
|
+
Returns:
|
|
55
|
+
sp.csr_matrix: The shape is (len(points), len(nn.X)).
|
|
56
|
+
"""
|
|
57
|
+
return nn.radius_neighbors_graph(
|
|
58
|
+
X=points,
|
|
59
|
+
radius=radius,
|
|
60
|
+
mode="distance" if return_distance else "connectivity",
|
|
61
|
+
)
|
|
62
|
+
|
|
63
|
+
@staticmethod
|
|
64
|
+
def __get_bool_batch(
|
|
65
|
+
numbers: np.ndarray,
|
|
66
|
+
batch: np.ndarray | list[int | bool] | None = None,
|
|
67
|
+
) -> tuple[np.ndarray, np.ndarray]:
|
|
68
|
+
"""Return numbers (int) & batch (bool) array."""
|
|
69
|
+
numbers = np.asarray(numbers, dtype=int)
|
|
70
|
+
if batch is None:
|
|
71
|
+
batch = np.ones_like(numbers, dtype=bool)
|
|
72
|
+
else:
|
|
73
|
+
batch = np.asarray(batch)
|
|
74
|
+
if batch.dtype not in (bool, np.bool_):
|
|
75
|
+
arr = np.zeros_like(numbers, dtype=bool)
|
|
76
|
+
arr[batch] = True
|
|
77
|
+
batch = arr
|
|
78
|
+
else:
|
|
79
|
+
assert batch.shape == numbers.shape
|
|
80
|
+
return numbers, batch
|
|
81
|
+
|
|
82
|
+
@classmethod
|
|
83
|
+
def get_distance_sparse_matrix(
|
|
84
|
+
cls,
|
|
85
|
+
p1: np.ndarray,
|
|
86
|
+
p2: np.ndarray | None = None,
|
|
87
|
+
cell: np.ndarray | Cell | None = None,
|
|
88
|
+
max_distance: float = float("inf"),
|
|
89
|
+
return_distance: bool = True,
|
|
90
|
+
) -> sp.csr_matrix:
|
|
91
|
+
"""Get Distance Matrix.
|
|
92
|
+
|
|
93
|
+
Args:
|
|
94
|
+
p1 (np.ndarray): The first group of positions.
|
|
95
|
+
p2 (np.ndarray | None, optional): The second group
|
|
96
|
+
of positions. Defaults to None. It means p2 is p1.
|
|
97
|
+
cell (np.ndarray, optional): The periodic cell. Defaults to None.
|
|
98
|
+
max_distance (float, optional): Defaults to float("inf").
|
|
99
|
+
return_distance (bool, optional): Defaults to True.
|
|
100
|
+
|
|
101
|
+
Returns:
|
|
102
|
+
coo_array: The Distance Matrix with shape (len(p1), len(p2)).
|
|
103
|
+
"""
|
|
104
|
+
if not isinstance(cell, Cell):
|
|
105
|
+
cell = Cell.new(cell=cell)
|
|
106
|
+
p1 = np.asarray(p1, dtype=float)
|
|
107
|
+
if p2 is None:
|
|
108
|
+
p2 = p1
|
|
109
|
+
else:
|
|
110
|
+
p2 = np.asarray(p2, dtype=float)
|
|
111
|
+
assert p1.ndim == p2.ndim == 2
|
|
112
|
+
assert p1.shape[1] == p2.shape[1] == 3
|
|
113
|
+
if len(p1) >= len(p2):
|
|
114
|
+
swap = False
|
|
115
|
+
else:
|
|
116
|
+
swap, p1, p2 = True, p2, p1
|
|
117
|
+
l1, l2 = p1.shape[0], p2.shape[0]
|
|
118
|
+
max_distance = float(max_distance)
|
|
119
|
+
assert max_distance > 0
|
|
120
|
+
|
|
121
|
+
distance = cls.__get_sparse_matrix(
|
|
122
|
+
nn=cls.__get_nn(X=p1),
|
|
123
|
+
points=p2,
|
|
124
|
+
radius=max_distance,
|
|
125
|
+
return_distance=bool(return_distance),
|
|
126
|
+
)
|
|
127
|
+
if cell.rank != 0:
|
|
128
|
+
distance = distance.astype(bool)
|
|
129
|
+
cell, _ = minkowski_reduce(cell)
|
|
130
|
+
assert isinstance(p2, np.ndarray)
|
|
131
|
+
p2 = wrap_positions(p2, cell, eps=0)
|
|
132
|
+
p1 = wrap_positions(p1, cell, eps=0)
|
|
133
|
+
assert max_distance != float("inf")
|
|
134
|
+
nn = cls.__get_nn(X=p1)
|
|
135
|
+
cell_pinv_T = np.transpose(pinv(cell))
|
|
136
|
+
n = [int(max_distance * norm(i)) + 1 for i in cell_pinv_T]
|
|
137
|
+
nrange3 = [range(-i, i + 1) for i in n]
|
|
138
|
+
for n1, n2, n3 in itertools.product(*nrange3):
|
|
139
|
+
if n1 <= 0 and (n2 < 0 or n2 == 0 and n3 < 0):
|
|
140
|
+
continue
|
|
141
|
+
shift_idxs = np.array([n1, n2, n3])
|
|
142
|
+
distance += cls.__get_sparse_matrix(
|
|
143
|
+
nn=nn,
|
|
144
|
+
radius=max_distance,
|
|
145
|
+
return_distance=False,
|
|
146
|
+
points=(p2 - shift_idxs @ cell), # type: ignore
|
|
147
|
+
)
|
|
148
|
+
distance += distance.T
|
|
149
|
+
if return_distance:
|
|
150
|
+
src, tgt = distance.tocoo().coords
|
|
151
|
+
v: np.ndarray = p2[src] - p1[tgt] # type: ignore
|
|
152
|
+
_, d = find_mic(v, cell=cell, pbc=True)
|
|
153
|
+
distance = sp.csr_matrix(
|
|
154
|
+
(d, (src, tgt)),
|
|
155
|
+
shape=distance.shape,
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
assert distance.shape == (l2, l1), f"{distance.shape} != {(l2, l1)}"
|
|
159
|
+
return distance if swap else distance.T
|
|
160
|
+
|
|
161
|
+
@classmethod
|
|
162
|
+
def get_distance_reduce_array(
|
|
163
|
+
cls,
|
|
164
|
+
p1: np.ndarray,
|
|
165
|
+
p2: np.ndarray | None = None,
|
|
166
|
+
cell: np.ndarray | Cell | None = None,
|
|
167
|
+
max_distance: float = float("inf"),
|
|
168
|
+
reduce_axis: Literal[0, 1] = 0,
|
|
169
|
+
) -> np.ndarray:
|
|
170
|
+
d_sp = cls.get_distance_sparse_matrix(p1, p2, cell, max_distance)
|
|
171
|
+
d_coo = sp.coo_matrix(d_sp, copy=False)
|
|
172
|
+
d_coo2 = sparse.COO(
|
|
173
|
+
coords=d_coo.coords,
|
|
174
|
+
fill_value=np.inf,
|
|
175
|
+
shape=d_coo.shape,
|
|
176
|
+
data=np.where(d_coo.data == 0, np.inf, d_coo.data),
|
|
177
|
+
)
|
|
178
|
+
return d_coo2.min(axis=reduce_axis).todense()
|
|
179
|
+
|
|
180
|
+
@classmethod
|
|
181
|
+
def get_adjacency_sparse_matrix(
|
|
182
|
+
cls,
|
|
183
|
+
numbers: np.ndarray,
|
|
184
|
+
geometry: np.ndarray,
|
|
185
|
+
batch: np.ndarray | list[int | bool] | None = None,
|
|
186
|
+
batch_other: np.ndarray | list[int | bool] | None = None,
|
|
187
|
+
cov_multiply_factor: float = 1.0,
|
|
188
|
+
cov_plus_factor: float = 0.5,
|
|
189
|
+
cell: Cell | None = None,
|
|
190
|
+
) -> sp.csr_matrix:
|
|
191
|
+
"""Get Adjacency Matrix by distance.
|
|
192
|
+
|
|
193
|
+
Args:
|
|
194
|
+
numbers (np.ndarray): the atomic number.
|
|
195
|
+
geometry (np.ndarray): the atomic geometry.
|
|
196
|
+
batch (np.ndarray | list[int | bool] | None):
|
|
197
|
+
the batch atoms of calculation. Default to None.
|
|
198
|
+
batch_other (np.ndarray | list[int | bool] | None):
|
|
199
|
+
the other batch atoms of calculation. Default to None.
|
|
200
|
+
cell (Cell, optional): the periodic cell. Defaults to None.
|
|
201
|
+
cov_multiply_factor (float): The multiply factor(1.0).
|
|
202
|
+
cov_plus_factor (float): The plus factor(0.5).
|
|
203
|
+
|
|
204
|
+
Note:
|
|
205
|
+
connected = bool(
|
|
206
|
+
distance < (ri + rj)
|
|
207
|
+
* cov_multiply_factor # Default to 1.0
|
|
208
|
+
+ cov_plus_factor means # Default to 0.5
|
|
209
|
+
)
|
|
210
|
+
|
|
211
|
+
Returns:
|
|
212
|
+
sp.csr_matrix: The Adjacency Matrix.
|
|
213
|
+
"""
|
|
214
|
+
_, batch_other = cls.__get_bool_batch(numbers, batch_other)
|
|
215
|
+
numbers, batch = cls.__get_bool_batch(numbers, batch)
|
|
216
|
+
geometry = np.asarray(geometry, dtype=float)
|
|
217
|
+
max_d = 2 * max(COV_R[numbers])
|
|
218
|
+
max_d *= cov_multiply_factor
|
|
219
|
+
max_d += cov_plus_factor
|
|
220
|
+
n = len(numbers)
|
|
221
|
+
|
|
222
|
+
distance = cls.get_distance_sparse_matrix(
|
|
223
|
+
p1=geometry[batch],
|
|
224
|
+
p2=geometry[batch_other],
|
|
225
|
+
max_distance=max_d,
|
|
226
|
+
return_distance=True,
|
|
227
|
+
cell=cell,
|
|
228
|
+
).tocoo()
|
|
229
|
+
d, (src, tgt) = distance.data, distance.coords
|
|
230
|
+
if distance.shape != (np.sum(batch), np.sum(batch_other)):
|
|
231
|
+
if distance.shape == (n, n):
|
|
232
|
+
tgt = np.arange(n)[batch_other][tgt]
|
|
233
|
+
src = np.arange(n)[batch][src]
|
|
234
|
+
# if error, check distance.shape, src & tgt
|
|
235
|
+
|
|
236
|
+
deq = COV_R[numbers[src]] + COV_R[numbers[tgt]]
|
|
237
|
+
deq = cov_multiply_factor * deq + cov_plus_factor
|
|
238
|
+
mask = np.logical_and(1e-3 < d, d < deq)
|
|
239
|
+
src, tgt = src[mask], tgt[mask]
|
|
240
|
+
d = np.ones_like(src, dtype=bool)
|
|
241
|
+
result = sp.csr_matrix(
|
|
242
|
+
(d, (src, tgt)),
|
|
243
|
+
shape=distance.shape,
|
|
244
|
+
dtype=bool,
|
|
245
|
+
)
|
|
246
|
+
assert result.dtype in (bool, np.bool_)
|
|
247
|
+
return result
|
|
248
|
+
|
|
249
|
+
@classmethod
|
|
250
|
+
def get_is_inner(
|
|
251
|
+
cls,
|
|
252
|
+
index: int,
|
|
253
|
+
numbers: np.ndarray,
|
|
254
|
+
geometry: np.ndarray,
|
|
255
|
+
adjacency_matrix: sp.csr_matrix | sp.spmatrix | sp.sparray,
|
|
256
|
+
cell: Cell | None = None,
|
|
257
|
+
) -> bool:
|
|
258
|
+
"""Check whether the given atom is inner by distance.
|
|
259
|
+
|
|
260
|
+
Args:
|
|
261
|
+
index (int): the index of the atom.
|
|
262
|
+
numbers (np.ndarray): the atomic number.
|
|
263
|
+
geometry (np.ndarray): the atomic geometry.
|
|
264
|
+
adjacency_matrix (sp.csr_matrix): the adjacency matrix.
|
|
265
|
+
cell (Cell, optional): the periodic cell. Defaults to None.
|
|
266
|
+
|
|
267
|
+
Returns:
|
|
268
|
+
bool: whether the given atom is inner.
|
|
269
|
+
"""
|
|
270
|
+
geometry = np.asarray(geometry, dtype=float)
|
|
271
|
+
adjacency_matrix = sp.csr_matrix(adjacency_matrix)
|
|
272
|
+
adjacency_matrix += adjacency_matrix.T
|
|
273
|
+
# get neighbors of the index atom
|
|
274
|
+
assert isinstance(adjacency_matrix, sp.csr_matrix)
|
|
275
|
+
nbr = sp.coo_array(adjacency_matrix[index, :])
|
|
276
|
+
idxs_neighbor = np.setdiff1d(nbr.coords[1], [index])
|
|
277
|
+
if len(idxs_neighbor) == 0:
|
|
278
|
+
return False
|
|
279
|
+
else:
|
|
280
|
+
# create sphere surface points sampling
|
|
281
|
+
mesh = inverse_3d_sphere_surface_sampling(1000)
|
|
282
|
+
mesh = geometry[index] + mesh * COV_R[numbers[index]]
|
|
283
|
+
# get distance between mesh and neighbor
|
|
284
|
+
distance = cls.get_distance_sparse_matrix(
|
|
285
|
+
p1=mesh,
|
|
286
|
+
p2=np.atleast_2d(geometry[idxs_neighbor]),
|
|
287
|
+
max_distance=(
|
|
288
|
+
np.max(COV_R[numbers[idxs_neighbor]])
|
|
289
|
+
+ float(COV_R[numbers[index]])
|
|
290
|
+
),
|
|
291
|
+
return_distance=True,
|
|
292
|
+
cell=cell,
|
|
293
|
+
).tocoo()
|
|
294
|
+
# get mesh mask array
|
|
295
|
+
d, (src, tgt) = distance.data, distance.coords
|
|
296
|
+
mask = d >= COV_R[tgt]
|
|
297
|
+
d, src, tgt = d[mask], src[mask], tgt[mask]
|
|
298
|
+
d, shape = np.ones_like(d, dtype=bool), distance.shape
|
|
299
|
+
mask = sp.csr_matrix(
|
|
300
|
+
(d, (src, tgt)),
|
|
301
|
+
shape=shape,
|
|
302
|
+
dtype=bool,
|
|
303
|
+
)
|
|
304
|
+
assert mask.shape == (len(mesh), len(idxs_neighbor))
|
|
305
|
+
mask_array = mask.sum(axis=1, dtype=bool)
|
|
306
|
+
# set is_inner
|
|
307
|
+
return bool(np.sum(mask_array) == len(mesh))
|
|
308
|
+
|
|
309
|
+
@classmethod
|
|
310
|
+
def get_is_inner_array(
|
|
311
|
+
cls,
|
|
312
|
+
numbers: np.ndarray,
|
|
313
|
+
geometry: np.ndarray,
|
|
314
|
+
adjacency_matrix: sp.csr_matrix | sp.spmatrix | sp.sparray,
|
|
315
|
+
batch: np.ndarray | list[int | bool] | None = None,
|
|
316
|
+
cell: Cell | None = None,
|
|
317
|
+
) -> list[bool] | np.ndarray:
|
|
318
|
+
"""Check whether each atom is inner by distance.
|
|
319
|
+
|
|
320
|
+
Args:
|
|
321
|
+
numbers (np.ndarray): the atomic number.
|
|
322
|
+
geometry (np.ndarray): the atomic geometry.
|
|
323
|
+
adjacency_matrix (sp.csr_matrix): the adjacency matrix.
|
|
324
|
+
batch (np.ndarray | list[int | bool], optional):
|
|
325
|
+
the batch atoms of calculation. Default to None.
|
|
326
|
+
cell (Cell, optional): the periodic cell. Defaults to None.
|
|
327
|
+
|
|
328
|
+
Returns:
|
|
329
|
+
list[bool]: whether each atom is inner.
|
|
330
|
+
"""
|
|
331
|
+
numbers, batch = cls.__get_bool_batch(numbers, batch)
|
|
332
|
+
result = np.zeros_like(batch, dtype=bool)
|
|
333
|
+
for index, in_batch in enumerate(batch):
|
|
334
|
+
if in_batch:
|
|
335
|
+
result[index] = cls.get_is_inner(
|
|
336
|
+
index=index,
|
|
337
|
+
numbers=numbers,
|
|
338
|
+
geometry=geometry,
|
|
339
|
+
adjacency_matrix=adjacency_matrix,
|
|
340
|
+
cell=cell,
|
|
341
|
+
)
|
|
342
|
+
return result
|
|
343
|
+
|
|
344
|
+
|
|
345
|
+
distance_factory = _DistanceFactory
|
|
346
|
+
get_distance_sparse_matrix = distance_factory.get_distance_sparse_matrix
|
|
347
|
+
get_adjacency_sparse_matrix = distance_factory.get_adjacency_sparse_matrix
|
|
348
|
+
get_is_inner_array = distance_factory.get_is_inner_array
|
|
349
|
+
|
|
350
|
+
|
|
351
|
+
def inverse_3d_sphere_surface_sampling(n: int) -> np.ndarray:
|
|
352
|
+
"""The 3D Sphere Surface Sampling Based on Inverse Transform.
|
|
353
|
+
|
|
354
|
+
Reference:
|
|
355
|
+
https://en.wikipedia.org/wiki/Inverse_transform_sampling
|
|
356
|
+
http://corysimon.github.io/articles/uniformdistn-on-sphere/
|
|
357
|
+
https://niepp.github.io/2021/12/09/uniform-sampling-on-sphere.html
|
|
358
|
+
|
|
359
|
+
Args:
|
|
360
|
+
n (int): the number of points for sampling.
|
|
361
|
+
|
|
362
|
+
Returns:
|
|
363
|
+
np.ndarray: (2n, 3) shape float numpy.ndarray for half sphere
|
|
364
|
+
"""
|
|
365
|
+
u = np.random.uniform(0, 1, size=n)
|
|
366
|
+
v = np.random.uniform(0, 1, size=n)
|
|
367
|
+
# inversion method
|
|
368
|
+
phi = v * 2 * np.pi
|
|
369
|
+
cos_theta = 1 - u
|
|
370
|
+
sin_theta = np.sqrt(1 - cos_theta * cos_theta)
|
|
371
|
+
x = np.cos(phi) * sin_theta
|
|
372
|
+
y = np.sin(phi) * sin_theta
|
|
373
|
+
z0, z1 = cos_theta, -cos_theta
|
|
374
|
+
half0 = np.column_stack([x, y, z0])
|
|
375
|
+
half1 = np.column_stack([x, y, z1])
|
|
376
|
+
return np.vstack([half0, half1])
|