def test_size_strategy(self):
bag = BaseNNGet(cutoff=5.0, nn_strategy="MinimumDistanceNNAll")
for i in self.data0_3:
center_indices, atom_nbr_idx, bond_states, bonds, center_prop = bag.convert(
i)
print(center_indices.shape)
print(atom_nbr_idx.shape)
print(bond_states.shape)
print(bonds.shape)
print(center_prop.shape)
print("next")
def test_size_radius(self):
bag = BaseNNGet(cutoff=5.0, nn_strategy="find_points_in_spheres")
for i in self.data0_3:
center_indices, atom_nbr_idx, bond_states, bonds, center_prop = bag.convert(
i)
print(center_indices.shape)
print(atom_nbr_idx.shape)
print(bond_states.shape)
print(bonds.shape)
print(center_prop.shape)
print("next")
def test_get1111(self):
bag = BaseNNGet(nn_strategy="CrystalNN")
for i in self.data0_checked:
resultt = bag.convert(i)
print(resultt)
def test_get911(self):
bag = BaseNNGet(nn_strategy="BrunnerNN_real")
for i in self.data0_checked:
resultt = bag.convert(i)
print(resultt)
def test_get199(self):
bag = BaseNNGet(cutoff=5.0, nn_strategy="find_points_in_spheres")
for i in self.data0_checked:
resultt = bag.convert(i)
print(resultt)
def test_get44(self):
bag = BaseNNGet(nn_strategy="MinimumDistanceNN")
for i in self.data0_checked:
resultt = bag.convert(i)
print(resultt)
def test_get10(self):
bag = BaseNNGet(nn_strategy="EconNN")
for i in self.data0_3:
resultt = bag.convert(i)
print(resultt)
def test_get9(self):
bag = BaseNNGet(nn_strategy="BrunnerNN_reciprocal")
for i in self.data0_3:
resultt = bag.convert(i)
print(resultt)
def test_get2(self):
bag = BaseNNGet(nn_strategy="UserVoronoiNN")
for i in self.data0_3:
resultt = bag.convert(i)
print(resultt)
def test_get(self):
bag = BaseNNGet(UserVoronoiNN)
for i in self.data0_3:
resultt = bag.convert(i)
print(resultt)
def __init__(self,
nn_strategy: Union[str,
NearNeighbors] = "find_points_in_spheres",
bond_generator: [BaseNNGet, BaseDesGet, str] = None,
atom_converter: Converter = None,
bond_converter: Converter = None,
state_converter: Converter = None,
return_bonds: str = "all",
cutoff: float = 5.0,
**kwargs):
"""
Parameters
----------
nn_strategy : str
NearNeighbor strategy.
For bond_converter ="BaseNNGet": ["BrunnerNN_reciprocal", "BrunnerNN_real", "BrunnerNN_relative",
"EconNN", "CrystalNN", "MinimumDistanceNNAll", "find_points_in_spheres","UserVoronoiNN"]
For bond_converter ="BaseDesGet": ["ACSF","BehlerParrinello","EAD","EAMD","SOAP","SO3","SO4_Bispectrum","wACSF"]
See Also:
``BaseNNGet`` :
:class:`featurebox.featurizers.envir.local_env.MinimumDistanceNNAll`,
``BaseDesGet`` :
:mod:`featurebox.featurizers.descriptors`,
:class:`featurebox.featurizers.descriptors.SOAP.SOAP`,
atom_converter: BinaryMap
atom features converter.
See Also:
:class:`featurebox.featurizers.atom.mapper.AtomTableMap` , :class:`featurebox.featurizers.atom.mapper.AtomJsonMap` ,
:class:`featurebox.featurizers.atom.mapper.AtomPymatgenPropMap`, :class:`featurebox.featurizers.atom.mapper.AtomTableMap`
bond_converter : Converter
bond features converter, default=None.
state_converter : Converter
state features converter.
See Also:
:class:`featurebox.featurizers.state.state_mapper.StructurePymatgenPropMap`
:mod:`featurebox.featurizers.state.statistics`
:mod:`featurebox.featurizers.state.union`
bond_generator : _BaseEnvGet, str
bond features converter.
The function of this, is to convert data format to a fixed format.
"BaseNNGet" or "BaseDesGet" or defined BaseNNGet,BaseDesGet object, default "BaseNNGet".
1, BaseDesGet or 2, BaseDesGet. or there name.
if object offered, rather str, the nn_strategy would use the nn_strategy in Converter.
See Also:
:class:`featurebox.featurizers.envir.environment.BaseNNGet` ,
:class:`featurebox.featurizers.envir.environment.BaseDesGet`
return_bonds: "all","bonds","bond_state"
which bond property return. default "all".
``"bonds_state"`` : bond properties and ``"bonds"`` : atoms number near this center atom.
cutoff: float
Whether to use depends on the ``nn_strategy``.
**kwargs:
"""
super().__init__(**kwargs)
self.return_bonds = return_bonds
self.cutoff = cutoff
if bond_generator is None: # default use NNDict
self.nn_strategy = get_marked_class(nn_strategy, NNDict)
# there use the universal parameter, custom it please
self.bond_generator = BaseNNGet(self.nn_strategy,
numerical_tol=1e-8,
pbc=None,
cutoff=self.cutoff)
elif isinstance(bond_generator, str): # new add "BaseDesGet"
self.nn_strategy = get_marked_class(nn_strategy,
env_method[bond_generator])
# there use the universal parameter, custom it please
self.bond_generator = env_names[bond_generator](self.nn_strategy,
self.cutoff,
numerical_tol=1e-8,
pbc=None,
cutoff=self.cutoff)
else: # defined BaseDesGet or BaseNNGet
self.bond_generator = bond_generator
self.nn_strategy = self.bond_generator.nn_strategy
self.atom_converter = atom_converter or self._get_dummy_converter()
self.bond_converter = bond_converter or self._get_dummy_converter()
self.state_converter = state_converter or self._get_dummy_converter()
self.graph_data_name = ["atom", "bond", "state", 'atom_nbr_idx']
self.cutoff = cutoff
class _StructureGraph(BaseFeature):
"""
Preferential use of _StructureGraphFixedRadius.
This is a base class for converting converting structure into graphs or model inputs.
Methods to be implemented are follows:
convert(self, structure)
This is to convert a structure into a graph dictionary
"""
def __init__(self,
nn_strategy: Union[str,
NearNeighbors] = "find_points_in_spheres",
bond_generator: [BaseNNGet, BaseDesGet, str] = None,
atom_converter: Converter = None,
bond_converter: Converter = None,
state_converter: Converter = None,
return_bonds: str = "all",
cutoff: float = 5.0,
**kwargs):
"""
Parameters
----------
nn_strategy : str
NearNeighbor strategy.
For bond_converter ="BaseNNGet": ["BrunnerNN_reciprocal", "BrunnerNN_real", "BrunnerNN_relative",
"EconNN", "CrystalNN", "MinimumDistanceNNAll", "find_points_in_spheres","UserVoronoiNN"]
For bond_converter ="BaseDesGet": ["ACSF","BehlerParrinello","EAD","EAMD","SOAP","SO3","SO4_Bispectrum","wACSF"]
See Also:
``BaseNNGet`` :
:class:`featurebox.featurizers.envir.local_env.MinimumDistanceNNAll`,
``BaseDesGet`` :
:mod:`featurebox.featurizers.descriptors`,
:class:`featurebox.featurizers.descriptors.SOAP.SOAP`,
atom_converter: BinaryMap
atom features converter.
See Also:
:class:`featurebox.featurizers.atom.mapper.AtomTableMap` , :class:`featurebox.featurizers.atom.mapper.AtomJsonMap` ,
:class:`featurebox.featurizers.atom.mapper.AtomPymatgenPropMap`, :class:`featurebox.featurizers.atom.mapper.AtomTableMap`
bond_converter : Converter
bond features converter, default=None.
state_converter : Converter
state features converter.
See Also:
:class:`featurebox.featurizers.state.state_mapper.StructurePymatgenPropMap`
:mod:`featurebox.featurizers.state.statistics`
:mod:`featurebox.featurizers.state.union`
bond_generator : _BaseEnvGet, str
bond features converter.
The function of this, is to convert data format to a fixed format.
"BaseNNGet" or "BaseDesGet" or defined BaseNNGet,BaseDesGet object, default "BaseNNGet".
1, BaseDesGet or 2, BaseDesGet. or there name.
if object offered, rather str, the nn_strategy would use the nn_strategy in Converter.
See Also:
:class:`featurebox.featurizers.envir.environment.BaseNNGet` ,
:class:`featurebox.featurizers.envir.environment.BaseDesGet`
return_bonds: "all","bonds","bond_state"
which bond property return. default "all".
``"bonds_state"`` : bond properties and ``"bonds"`` : atoms number near this center atom.
cutoff: float
Whether to use depends on the ``nn_strategy``.
**kwargs:
"""
super().__init__(**kwargs)
self.return_bonds = return_bonds
self.cutoff = cutoff
if bond_generator is None: # default use NNDict
self.nn_strategy = get_marked_class(nn_strategy, NNDict)
# there use the universal parameter, custom it please
self.bond_generator = BaseNNGet(self.nn_strategy,
numerical_tol=1e-8,
pbc=None,
cutoff=self.cutoff)
elif isinstance(bond_generator, str): # new add "BaseDesGet"
self.nn_strategy = get_marked_class(nn_strategy,
env_method[bond_generator])
# there use the universal parameter, custom it please
self.bond_generator = env_names[bond_generator](self.nn_strategy,
self.cutoff,
numerical_tol=1e-8,
pbc=None,
cutoff=self.cutoff)
else: # defined BaseDesGet or BaseNNGet
self.bond_generator = bond_generator
self.nn_strategy = self.bond_generator.nn_strategy
self.atom_converter = atom_converter or self._get_dummy_converter()
self.bond_converter = bond_converter or self._get_dummy_converter()
self.state_converter = state_converter or self._get_dummy_converter()
self.graph_data_name = ["atom", "bond", "state", 'atom_nbr_idx']
self.cutoff = cutoff
def __add__(self, other):
raise TypeError("There is no add.")
def convert(self,
structure: Structure,
state_attributes: List = None) -> Dict:
"""
Take a pymatgen structure and convert it to a index-type graph representation
The graph will have node, distance, index1, index2, where node is a vector of Z number
of atoms in the structure, index1 and index2 mark the atom indices forming the bond and separated by
distance.
For state attributes, you can set structure.state = [[xx, xx]] beforehand or the algorithm would
take default [[0, 0]]
Parameters
----------
state_attributes: list
state attributes
structure: Structure
Returns
-------
``atom_fea``: np.ndarray, shape (N, atom_fea_len)
center properties.
``nbr_fea``: np.ndarray, shape (N, fill_size, atom_fea_len).
neighbor_indexes for each center_index.
`fill_size` default = 5.
``state_fea``: np.ndarray, shape (state_fea_len,)
state feature.
``atom_nbr_idx``: np.ndarray, shape (N, fill_size)
neighbor for each center, fill_size default is 5.
"""
if state_attributes is not None:
state_attributes = np.array(state_attributes)
else:
state_attributes = np.array([0.0, 0.0], dtype="float32")
if isinstance(self.state_converter, DummyConverter):
pass
else:
state_attributes = np.concatenate(
(state_attributes,
np.array(self.state_converter.convert(structure)).ravel()))
center_indices, atom_nbr_idx, bond_states, bonds, center_prop = self.get_bond_fea(
structure)
if self.return_bonds == "all":
bondss = np.concatenate(
(bonds,
bond_states), axis=-1) if bonds is not None else bond_states
elif self.return_bonds == "bonds":
if bonds is not None:
bondss = bonds
else:
raise TypeError
elif self.return_bonds == "bonds_state":
bondss = bond_states
else:
raise NotImplementedError()
if isinstance(self.atom_converter, DummyConverter):
atoms_numbers = np.array(
structure.atomic_numbers)[center_indices].reshape(-1, 1)
atoms = self.atom_converter.convert(atoms_numbers)
elif isinstance(self.atom_converter, ConverterCat):
self.atom_converter.force_concatenate = True # just accept the data could be concatenate as one array.
atoms = self.atom_converter.convert(structure)
atoms = np.array([atoms[round(i)] for i in center_indices])
else:
atoms = self.atom_converter.convert(structure)
atoms = np.array([atoms[round(i)] for i in center_indices])
bonds = self.bond_converter.convert(bondss)
# atoms number in the first column.
if atoms.shape[1] == 1:
if center_prop.shape[1] > 1:
atoms = np.concatenate((np.array(atoms), center_prop), axis=1)
else:
pass
elif atoms.shape[1] > 1:
atoms_numbers = np.array(
structure.atomic_numbers)[center_indices].reshape(-1, 1)
if center_prop.shape[1] > 1:
atoms = np.concatenate(
(atoms_numbers, np.array(atoms), center_prop), axis=1)
else:
atoms = np.concatenate((atoms_numbers, np.array(atoms)),
axis=1)
else:
raise TypeError(
"Bad Converter for: atoms = self.atom_converter.convert(atoms_numbers)"
)
return {
"atom": atoms,
"bond": bonds,
"state": state_attributes,
"atom_nbr_idx": atom_nbr_idx
}
def get_bond_fea(self, structure: Structure):
"""
Get atom features from structure, may be overwritten.
"""
# assert hasattr(self.bond_generator, "convert")
return self.bond_generator.convert(structure)
def __call__(self, structure: Structure, *args, **kwargs) -> Dict:
return self.convert(structure, *args, **kwargs)
@staticmethod
def _get_dummy_converter() -> DummyConverter:
return DummyConverter()
def _transform(self,
structures: List[Structure],
state_attributes: List = None):
"""
Parameters
----------
structures:list
preprocessing of samples need to transform to Graph.
state_attributes:List
preprocessing of samples need to add to Graph.
Returns
-------
list of graphs:
List of dict
"""
if state_attributes is None:
state_attributes = [None] * len(structures)
assert isinstance(structures, Iterable)
if hasattr(structures, "__len__"):
assert len(structures) > 0, "Empty input data!"
iterables = zip(structures, state_attributes)
if not self.batch_calculate:
rets = parallelize(self.n_jobs,
self._wrapper,
iterables,
tq=True,
respective=True)
ret, self.support_ = zip(*rets)
else:
rets = batch_parallelize(self.n_jobs,
self._wrapper,
iterables,
respective=True,
tq=True,
batch_size=self.batch_size)
ret, self.support_ = zip(*rets)
return ret
def get_flat_data(self, graphs: List[Dict]) -> tuple:
"""
Expand the graph dictionary to form a list of features and targets tensors.
This is useful when the model is trained on assembled graphs on the fly.
Aim to get input for GraphGenerator.
Parameters
----------
graphs: list of dict
list of graph dictionary for each structure
Returns
-------
tuple(node_features, edges_features, stats_values, atom_nbr_idx , ***)
"""
output = [] # Will be a list of arrays
# Convert the graphs to matrices
for n, fi in enumerate(self.graph_data_name):
output.append([
np.array(gi[fi]) if isinstance(gi, dict) else np.array(gi[n])
for gi in graphs
])
return tuple(output)
def transform(self,
structures: List[Structure],
state_attributes: List = None):
"""
Parameters
----------
structures:list
preprocessing of samples need to transform to Graph.
state_attributes
preprocessing of samples need to add to Graph.
Returns
-------
``atom_fea``: list of np.ndarray, shape (N, atom_fea_len)
center properties.
``nbr_fea``: list of np.ndarray, shape (N, fill_size, atom_fea_len).
neighbor_indexes for each center_index.
`fill_size` default = 5.
``state_fea``: list of np.ndarray, shape (state_fea_len,)
state feature.
``atom_nbr_idx``: list of np.ndarray, shape (N, fill_size)
neighbor for each center, fill_size default is 5.
"""
return self.get_flat_data(self._transform(structures,
state_attributes))