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.
Files changed (39) hide show
  1. graphatoms-0.0.2/GraphAtoms/__init__.py +8 -0
  2. graphatoms-0.0.2/GraphAtoms/common/__init__.py +4 -0
  3. graphatoms-0.0.2/GraphAtoms/common/abc.py +285 -0
  4. graphatoms-0.0.2/GraphAtoms/common/error.py +16 -0
  5. graphatoms-0.0.2/GraphAtoms/common/geometry.py +376 -0
  6. graphatoms-0.0.2/GraphAtoms/common/ndarray.py +3 -0
  7. graphatoms-0.0.2/GraphAtoms/common/ndarray.pyi +5 -0
  8. graphatoms-0.0.2/GraphAtoms/common/rotation.py +101 -0
  9. graphatoms-0.0.2/GraphAtoms/common/string.py +33 -0
  10. graphatoms-0.0.2/GraphAtoms/containner/__init__.py +20 -0
  11. graphatoms-0.0.2/GraphAtoms/containner/_atomic.py +366 -0
  12. graphatoms-0.0.2/GraphAtoms/containner/_bonded.py +90 -0
  13. graphatoms-0.0.2/GraphAtoms/containner/_graph.py +704 -0
  14. graphatoms-0.0.2/GraphAtoms/containner/_mConnComp.py +45 -0
  15. graphatoms-0.0.2/GraphAtoms/containner/_mDataPyG.py +87 -0
  16. graphatoms-0.0.2/GraphAtoms/containner/_mFreeE.py +117 -0
  17. graphatoms-0.0.2/GraphAtoms/containner/_mIGraph.py +269 -0
  18. graphatoms-0.0.2/GraphAtoms/containner/_mNetworKit.py +12 -0
  19. graphatoms-0.0.2/GraphAtoms/containner/_mNetworkX.py +12 -0
  20. graphatoms-0.0.2/GraphAtoms/containner/_mRDKit.py +337 -0
  21. graphatoms-0.0.2/GraphAtoms/containner/_mRustworkX.py +12 -0
  22. graphatoms-0.0.2/GraphAtoms/containner/_sysCluster.py +121 -0
  23. graphatoms-0.0.2/GraphAtoms/containner/_sysGas.py +81 -0
  24. graphatoms-0.0.2/GraphAtoms/containner/_system.py +15 -0
  25. {graphatoms-0.0.1 → graphatoms-0.0.2}/GraphAtoms.egg-info/PKG-INFO +1 -2
  26. graphatoms-0.0.2/GraphAtoms.egg-info/SOURCES.txt +36 -0
  27. {graphatoms-0.0.1 → graphatoms-0.0.2}/GraphAtoms.egg-info/requires.txt +0 -1
  28. {graphatoms-0.0.1 → graphatoms-0.0.2}/PKG-INFO +1 -2
  29. {graphatoms-0.0.1 → graphatoms-0.0.2}/pyproject.toml +1 -7
  30. graphatoms-0.0.1/GraphAtoms.egg-info/SOURCES.txt +0 -12
  31. {graphatoms-0.0.1 → graphatoms-0.0.2}/GraphAtoms.egg-info/dependency_links.txt +0 -0
  32. {graphatoms-0.0.1 → graphatoms-0.0.2}/GraphAtoms.egg-info/top_level.txt +0 -0
  33. {graphatoms-0.0.1 → graphatoms-0.0.2}/LICENSE +0 -0
  34. {graphatoms-0.0.1 → graphatoms-0.0.2}/README.md +0 -0
  35. {graphatoms-0.0.1 → graphatoms-0.0.2}/setup.cfg +0 -0
  36. {graphatoms-0.0.1 → graphatoms-0.0.2}/tests/test_abc.py +0 -0
  37. {graphatoms-0.0.1 → graphatoms-0.0.2}/tests/test_container.py +0 -0
  38. {graphatoms-0.0.1 → graphatoms-0.0.2}/tests/test_mRDKit.py +0 -0
  39. {graphatoms-0.0.1 → graphatoms-0.0.2}/tests/test_utils.py +0 -0
@@ -0,0 +1,8 @@
1
+ # ruff: noqa
2
+ from .containner import Cluster, Gas, System
3
+
4
+ __all__ = [
5
+ "Cluster",
6
+ "Gas",
7
+ "System",
8
+ ]
@@ -0,0 +1,4 @@
1
+ """The module for common use.
2
+
3
+ Many other modules depend on this module.
4
+ """
@@ -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])
@@ -0,0 +1,3 @@
1
+ from numpydantic import NDArray, Shape # noqa: D100
2
+
3
+ __all__ = ["NDArray", "Shape"]
@@ -0,0 +1,5 @@
1
+ from numpy import ndarray as Numpyndarray
2
+ from numpydantic import Shape
3
+
4
+ NDArray = Numpyndarray
5
+ __all__ = ["Shape", "NDArray"]