amd.io.periodicset_from_pymatgen_cifblock()

Complexity

Conditions

28

Size

Total Lines	173
Code Lines	96

Duplication

Lines	24
Ratio	13.87 %

Importance

Changes

0

"""Tools for reading crystals from files, or from the CSD with
def periodicset_from_pymatgen_cifblock(
        block,
        remove_hydrogens: bool = False,
        disorder: str = 'skip'
) -> PeriodicSet:
    """Convert a :class:`pymatgen.io.cif.CifBlock` object to a
    :class:`amd.PeriodicSet <.periodicset.PeriodicSet>`.
    :class:`pymatgen.io.cif.CifBlock` is the type returned by
    :class:`pymatgen.io.cif.CifFile`.

    Parameters
    ----------
    block : :class:`pymatgen.io.cif.CifBlock`
        A pymatgen CifBlock object representing a crystal.
    remove_hydrogens : bool, optional
        Remove Hydrogens from the crystal.
    disorder : str, optional
        Controls how disordered structures are handled. Default is
        ``skip`` which skips any crystal with disorder, since disorder
        conflicts with the periodic set model. To read disordered
        structures anyway, choose either :code:`ordered_sites` to remove
        atoms with disorder or :code:`all_sites` include all atoms
        regardless of disorder.

    Returns
    -------
    :class:`amd.PeriodicSet <.periodicset.PeriodicSet>`
        Represents the crystal as a periodic set, consisting of a finite
        set of points (motif) and lattice (unit cell). Contains other
        useful data, e.g. the crystal's name and information about the
        asymmetric unit for calculation.

    Raises
    ------
    ParseError
        Raised if the structure can/should not be parsed for the
        following reasons: 1. No sites found or motif is empty after
        removing Hydrogens & disorder, 2. A site has missing
        coordinates, 3. :code:``disorder == 'skip'`` and disorder is
        found on any atom.
    """

    from pymatgen.io.cif import str2float

    odict = block.data

    # Unit cell
    cellpar = [odict.get(tag) for tag in _CIF_TAGS['cellpar']]
    if any(par in (None, '?', '.') for par in cellpar):
        raise ParseError(f'{block.header} has missing cell data')
    cell = cellpar_to_cell(np.array([str2float(v) for v in cellpar]))

    # Asymmetric unit coordinates
    asym_unit = [odict.get(tag) for tag in _CIF_TAGS['atom_site_fract']]
    # check for . and ?
    if None in asym_unit:
        asym_unit = [odict.get(tag) for tag in _CIF_TAGS['atom_site_cartn']]
        if None in asym_unit:
            raise ParseError(f'{block.header} has missing coordinates')
        else:
            raise ParseError(
                f'{block.header} uses _atom_site_Cartn_ tags for coordinates, '
                'only _atom_site_fract_ is supported'
            )
    asym_unit = list(zip(*asym_unit))
    asym_unit = [[str2float(coord) for coord in xyz] for xyz in asym_unit]

    # Labels
    asym_labels = odict.get('_atom_site_label')
    if asym_labels is None:
        asym_labels = [''] * len(asym_unit)

    # Atomic types
    asym_symbols = odict.get('_atom_site_type_symbol')
    if asym_symbols is not None:

This code seems to be duplicated in your project.

        asym_symbols_ = []
        for label in asym_symbols:
            sym = ''
            if label and label not in ('.', '?'):
                match = re.search(r'([A-Z][a-z]?)', label)
                if match is not None:
                    sym = match.group() 
            asym_symbols_.append(sym)
    else:
        asym_symbols_ = [''] * len(asym_unit)
    asym_types = [_ATOMIC_NUMBERS[s] for s in asym_symbols_]

    # Find where sites have disorder if necassary
    has_disorder = []
    if disorder != 'all_sites':
        occupancies = odict.get('_atom_site_occupancy')
        if occupancies is None:
            occupancies = np.ones((len(asym_unit), ))
        else:
            occupancies = np.array([str2float(occ) for occ in occupancies])
        labels = odict.get('_atom_site_label')
        if labels is None:
            labels = [''] * len(asym_unit)
        for lab, occ in zip(labels, occupancies):
            has_disorder.append(_has_disorder(lab, occ))

    # Remove sites with ?, . or other invalid string for coordinates
    invalid = []
    for i, xyz in enumerate(asym_unit):
        if not all(isinstance(coord, (int, float)) for coord in xyz):
            invalid.append(i)

    if invalid:
        warnings.warn('atoms without sites or missing data will be removed')
        asym_unit = [c for i, c in enumerate(asym_unit) if i not in invalid]
        asym_types = [c for i, c in enumerate(asym_types) if i not in invalid]
        if disorder != 'all_sites':
            has_disorder = [
                d for i, d in enumerate(has_disorder) if i not in invalid
            ]

    remove_sites = []

    if remove_hydrogens:
        remove_sites.extend((i for i, n in enumerate(asym_types) if n == 1))

The variable i does not seem to be defined in case the for loop on line 950 is not entered. Are you sure this can never be the case?

    # Remove atoms with fractional occupancy or raise ParseError
    if disorder != 'all_sites':

This code seems to be duplicated in your project.

        for i, dis in enumerate(has_disorder):
            if i in remove_sites:
                continue
            if dis:
                if disorder == 'skip':
                    raise ParseError(
                        f'{block.header} has disorder, pass '
                        "disorder='ordered_sites' or 'all_sites' to "
                        'remove/ignore disorder'
                    )
                elif disorder == 'ordered_sites':
                    remove_sites.append(i)

    # Asymmetric unit
    asym_unit = [c for i, c in enumerate(asym_unit) if i not in remove_sites]
    asym_types = [t for i, t in enumerate(asym_types) if i not in remove_sites]
    if len(asym_unit) == 0:
        raise ParseError(f'{block.header} has no valid sites')
    asym_unit = np.mod(np.array(asym_unit), 1)

    # recommended by pymatgen
    # asym_unit = _snap_small_prec_coords(asym_unit, 1e-4)

    # Remove overlapping sites unless disorder == 'all_sites'
    if disorder != 'all_sites':
        keep_sites = _unique_sites(asym_unit, _EQ_SITE_TOL)
        if not np.all(keep_sites):
            warnings.warn(
                'may have overlapping sites; duplicates will be removed'
            )
        asym_unit = asym_unit[keep_sites]
        asym_types = [sym for sym, keep in zip(asym_types, keep_sites) if keep]

    # Apply symmetries to asymmetric unit
    rot, trans = _get_syms_pymatgen(odict)
    frac_motif, invs = _expand_asym_unit(asym_unit, rot, trans, _EQ_SITE_TOL)
    _, wyc_muls = np.unique(invs, return_counts=True)
    wyc_muls = wyc_muls.astype(np.int32)
    asym_inds = np.zeros_like(wyc_muls, dtype=np.int32)
    asym_inds[1:] = np.cumsum(wyc_muls)[:-1]
    types = np.array([asym_types[i] for i in invs], dtype=np.uint8)
    motif = frac_motif @ cell

    return PeriodicSet(
        motif=motif,
        cell=cell,
        name=block.header,
        asymmetric_unit=asym_inds,
        wyckoff_multiplicities=wyc_muls,
        types=types
    )