How to use the megnet.data.graph.StructureGraph function in megnet

"""

    def __init__(self,
                 nn_strategy='MinimumDistanceNNAll',
                 atom_converter=None,
                 bond_converter=None,
                 cutoff=4.0
                 ):
        if bond_converter is None:
            bond_converter = GaussianDistance(np.linspace(0, 5, 100), 0.5)
        self.cutoff = cutoff
        super().__init__(nn_strategy=nn_strategy, atom_converter=atom_converter,
                         bond_converter=bond_converter, cutoff=self.cutoff)


class CrystalGraphWithBondTypes(StructureGraph):
    """
    Overwrite the bond attributes with bond types, defined simply by
    the metallicity of the atoms forming the bond. Three types of
    scenario is considered, nonmetal-nonmetal (type 0), metal-nonmetal (type 1), and
    metal-metal (type 2)

    """

    def __init__(self,
                 nn_strategy='VoronoiNN',
                 atom_converter=None,
                 bond_converter=None):
        super().__init__(nn_strategy=nn_strategy, atom_converter=atom_converter,
                         bond_converter=bond_converter)

    def convert(self, structure, state_attributes=None):

from megnet.data.graph import StructureGraph
import numpy as np
from megnet.data.graph import GaussianDistance
from monty.serialization import loadfn
from pathlib import Path
from copy import deepcopy
from pymatgen import Element

MODULE_DIR = Path(__file__).parent.absolute()


class CrystalGraph(StructureGraph):
    """
    Convert a crystal into a graph with z as atomic feature and distance as bond feature
    one can optionally include state features
    """

    def __init__(self,
                 nn_strategy='MinimumDistanceNNAll',
                 atom_converter=None,
                 bond_converter=None,
                 cutoff=4.0
                 ):
        if bond_converter is None:
            bond_converter = GaussianDistance(np.linspace(0, 5, 100), 0.5)
        self.cutoff = cutoff
        super().__init__(nn_strategy=nn_strategy, atom_converter=atom_converter,
                         bond_converter=bond_converter, cutoff=self.cutoff)

Chem = None

__date__ = '12/01/2018'

# List of features to use by default for each atom
_ATOM_FEATURES = ['element', 'chirality', 'formal_charge', 'ring_sizes',
                  'hybridization', 'donor', 'acceptor', 'aromatic']

# List of features to use by default for each bond
_BOND_FEATURES = ['bond_type', 'same_ring', 'spatial_distance', 'graph_distance']

# List of elements in library to use by default
_ELEMENTS = ['H', 'C', 'N', 'O', 'F']


class SimpleMolGraph(StructureGraph):
    """
    Default using all atom pairs as bonds. The distance between atoms are used
    as bond features. By default the distance is expanded using a Gaussian
    expansion with centers at np.linspace(0, 4, 20) and width of 0.5
    """
    def __init__(self,
                 nn_strategy='AllAtomPairs',
                 atom_converter=None,
                 bond_converter=None
                 ):
        if bond_converter is None:
            bond_converter = GaussianDistance(np.linspace(0, 4, 20), 0.5)
        super().__init__(nn_strategy=nn_strategy, atom_converter=atom_converter,
                         bond_converter=bond_converter)

if 'nn_strategy' in all_dict:
            nn_strategy = all_dict.pop('nn_strategy')
            all_dict.update({'nn_strategy': local_env.serialize(nn_strategy)})
        return all_dict

    @classmethod
    def from_dict(cls, d: Dict) -> 'StructureGraph':
        if 'nn_strategy' in d:
            nn_strategy = d.pop('nn_strategy')
            nn_strategy_obj = local_env.deserialize(nn_strategy)
            d.update({'nn_strategy': nn_strategy_obj})
            return super().from_dict(d)
        return super().from_dict(d)


class StructureGraphFixedRadius(StructureGraph):
    """
    This one uses a short cut to call find_points_in_spheres cython function in
    pymatgen. It is orders of magnitude faster than previous implementations
    """

    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]]

        Args:
            state_attributes: (list) state attributes

Default using all atom pairs as bonds. The distance between atoms are used
    as bond features. By default the distance is expanded using a Gaussian
    expansion with centers at np.linspace(0, 4, 20) and width of 0.5
    """
    def __init__(self,
                 nn_strategy='AllAtomPairs',
                 atom_converter=None,
                 bond_converter=None
                 ):
        if bond_converter is None:
            bond_converter = GaussianDistance(np.linspace(0, 4, 20), 0.5)
        super().__init__(nn_strategy=nn_strategy, atom_converter=atom_converter,
                         bond_converter=bond_converter)


class MolecularGraph(StructureGraph):
    """Class for generating the graph inputs from a molecule

    Computes many different features for the atoms and bonds in a molecule, and prepares them
    in a form compatible with MEGNet models. The :meth:`convert` method takes a OpenBabel molecule
    and, besides computing features, also encodes them in a form compatible with machine learning.
    Namely, the `convert` method one-hot encodes categorical variables and concatenates
    the atomic features

    ## Atomic Features

    This class can compute the following features for each atom

    - `atomic_num`: The atomic number
    - `element`: (categorical) Element identity. (Unlike `atomic_num`, element is one-hot-encoded)
    - `chirality`: (categorical) R, S, or not a Chiral center (one-hot encoded).
    - `formal_charge`: Formal charge of the atom