amd.io.periodicset_from_ase_cifblock()

Complexity

Conditions

31

Size

Total Lines	182
Code Lines	103

Duplication

Lines	24
Ratio	13.19 %

Importance

Changes

0

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

    Parameters
    ----------
    block : :class:`ase.io.cif.CIFBlock`
        An ase :class:`ase.io.cif.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 fails to be parsed for any of the
        following: 1. Required data is missing (e.g. cell parameters),
        2. The motif is empty after removing H or disordered sites,
        3. :code:``disorder == 'skip'`` and disorder is found on any
        atom.
    """

    import ase
    import ase.spacegroup

    # Unit cell
    cellpar = [block.get(tag) for tag in _CIF_TAGS['cellpar']]
    if None in cellpar:
        raise ParseError(f'{block.name} has missing cell data')
    cell = cellpar_to_cell(np.array(cellpar))

    # Asymmetric unit coordinates. ase removes uncertainty brackets
    asym_unit = [block.get(tag) for tag in _CIF_TAGS['atom_site_fract']]
    if None in asym_unit:
        asym_unit = [
            block.get(tag.lower()) for tag in _CIF_TAGS['atom_site_cartn']
        ]
        if None in asym_unit:
            raise ParseError(f'{block.name} has missing coordinates')
        else:
            raise ParseError(
                f'{block.name} uses _atom_site_Cartn_ tags for coordinates, '
                'only _atom_site_fract_ is supported'
            )
    asym_unit = list(zip(*asym_unit))

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

    # Atomic types
    asym_symbols = block.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 = block.get('_atom_site_occupancy')
        if occupancies is None:
            occupancies = [1] * len(asym_unit)
        for lab, occ in zip(asym_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 = [t for i, t 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, num in enumerate(asym_types) if num == 1)

    # 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.name} 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.name} 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 = [t for t, keep in zip(asym_types, keep_sites) if keep]

    # Get symmetry operations
    sitesym = block._get_any(_CIF_TAGS['symop'])
    if sitesym is None:
        label_or_num = block._get_any(
            [s.lower() for s in _CIF_TAGS['spacegroup_name']]
        )
        if label_or_num is None:
            label_or_num = block._get_any(
                [s.lower() for s in _CIF_TAGS['spacegroup_number']]
            )
        if label_or_num is None:
            warnings.warn('no symmetry data found, defaulting to P1')
            label_or_num = 1
        spg = ase.spacegroup.Spacegroup(label_or_num)
        rot, trans = spg.get_op()
    else:
        if isinstance(sitesym, str):
            sitesym = [sitesym]
        rot, trans = _parse_sitesyms(sitesym)

    frac_motif, invs = _expand_asym_unit(asym_unit, rot, trans, _EQ_SITE_TOL)
    _, wyc_muls = np.unique(invs, return_counts=True)
    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 = np.matmul(frac_motif, cell)

    return PeriodicSet(
        motif=motif,
        cell=cell,
        name=block.name,
        asym_unit=asym_inds,
        multiplicities=wyc_muls,
        types=types
    )