crystal

def __init__(
    self,
    unit_cell: UnitCell,
    space_group: SpaceGroup,
    asymmetric_unit: AsymmetricUnit,
    **kwargs,
):
    """
    Construct a new crystal.


    Arguments:
        unit_cell: The unit cell for this crystal i.e. the
            translational symmetry of the crystal structure.
        space_group: The space group symmetry of this crystal
            i.e. the generators for populating the unit cell given the
            asymmetric unit.
        asymmetric_unit: The asymmetric unit of this crystal.
             The sites of this combined with the space group will generate all
             translationally equivalent positions.
        **kwargs: Optional properties to (will populate the properties member) store
            about the the crystal structure.
    """

    self.space_group = space_group
    self.unit_cell = unit_cell
    self.asymmetric_unit = asymmetric_unit
    self.properties = {}
    self.properties.update(kwargs)

def as_P1(self) -> "Crystal":
    """Create a copy of this crystal in space group P 1, with the new
    asymmetric_unit consisting of self.unit_cell_molecules()"""
    return self.as_P1_supercell((1, 1, 1))

def as_P1_supercell(self, size) -> "Crystal":
    """
    Create a supercell of this crystal in space group P 1.

    Args:
        size (Tuple[int]): size of the P 1 supercell to be created

    Returns:
        Crystal object of a supercell in space group P 1
    """
    import itertools as it

    umax, vmax, wmax = size
    a, b, c = self.unit_cell.lengths
    sc = UnitCell.from_lengths_and_angles(
        (umax * a, vmax * b, wmax * c), self.unit_cell.angles
    )

    u = np.arange(umax)
    v = np.arange(vmax)
    w = np.arange(wmax)
    sc_mols = []
    for q, r, s in it.product(u, v, w):
        for uc_mol in self.unit_cell_molecules():
            sc_mols.append(
                uc_mol.translated(np.asarray([q, r, s]) @ self.unit_cell.lattice)
            )

    asym_pos = np.vstack([x.positions for x in sc_mols])
    asym_nums = np.hstack([x.atomic_numbers for x in sc_mols])
    asymmetric_unit = AsymmetricUnit(
        [Element[x] for x in asym_nums], sc.to_fractional(asym_pos)
    )
    new_crystal = Crystal(sc, SpaceGroup(1), asymmetric_unit)
    new_crystal.properties["titl"] = self.titl + "-P1-{}-{}-{}".format(*size)
    return new_crystal

def asymmetric_unit_partial_charges(self) -> np.ndarray:
    """
    Calculate the partial charges for the asymmetric unit of this
    crystal using the EEM method.

    Returns:
        an `ndarray` of atomic partial charges.
    """
    mols = self.symmetry_unique_molecules()
    charges = np.empty(len(self.asymmetric_unit), dtype=np.float32)
    for mol in mols:
        for idx, charge in zip(
            mol.properties["asymmetric_unit_atoms"], mol.partial_charges
        ):
            charges[idx] = charge
    return charges

def atom_group_shape_descriptors(self, atoms, l_max=5, radius=6.0) -> np.ndarray:
    """Calculate the shape descriptors[1,2] for the given atomic
    group in this crystal.

    Args:
        atoms (Tuple): atoms to include in the as the 'inside' of the shape description.
        l_max (int, optional): maximum level of angular momenta to include in the spherical harmonic
            transform of the molecular shape function.
        radius (float, optional): maximum distance (Angstroms) to include surroundings
            in the shape description

    Returns:
        shape description vector

    References:
    ```
    [1] PR Spackman et al. Sci. Rep. 6, 22204 (2016)
        https://dx.doi.org/10.1038/srep22204
    [2] PR Spackman et al. Angew. Chem. 58 (47), 16780-16784 (2019)
        https://dx.doi.org/10.1002/anie.201906602
    ```
    """
    from chmpy.shape import SHT, stockholder_weight_descriptor

    sph = SHT(l_max=l_max)
    inside, outside = self.atom_group_surroundings(atoms, radius=radius)
    m = Molecule.from_arrays(*inside)
    c = np.array(m.centroid, dtype=np.float32)
    dists = np.linalg.norm(m.positions - c, axis=1)
    bounds = np.min(dists) / 2, np.max(dists) + 10.0
    return np.asarray(
        stockholder_weight_descriptor(
            sph, *inside, *outside, origin=c, bounds=bounds
        )
    )

def atom_group_surroundings(self, atoms, radius=6.0) -> Tuple:
    """
    Calculate all atoms within the given `radius` of the specified
    group of atoms in the asymetric unit.

    Arguments:
        radius (float): the maximum distance (Angstroms) from the origin for inclusion

    Returns:
        A list of atomic number, Cartesian position for both the
        atomic sites in question and their surroundings (as an array)
    """
    hklmax = np.array([-np.inf, -np.inf, -np.inf])
    hklmin = np.array([np.inf, np.inf, np.inf])
    frac_radius = radius / np.array(self.unit_cell.lengths)
    mol = self.symmetry_unique_molecules()[0]
    central_positions = self.to_fractional(mol.positions[atoms])
    central_elements = mol.atomic_numbers[atoms]
    central_cart_positions = mol.positions[atoms]

    for pos in central_positions:
        hklmax = np.maximum(hklmax, np.ceil(frac_radius + pos))
        hklmin = np.minimum(hklmin, np.floor(pos - frac_radius))
    hmax, kmax, lmax = hklmax.astype(int)
    hmin, kmin, lmin = hklmin.astype(int)
    slab = self.slab(bounds=((hmin, kmin, lmin), (hmax, kmax, lmax)))
    elements = slab["element"]
    positions = slab["cart_pos"]
    tree = KDTree(positions)
    keep = np.zeros(positions.shape[0], dtype=bool)

    this_mol = []
    for pos in central_cart_positions:
        idxs = tree.query_ball_point(pos, radius)
        d, nn = tree.query(pos)
        keep[idxs] = True
        if d < 1e-3:
            this_mol.append(nn)
            keep[this_mol] = False
    return (
        (central_elements, central_cart_positions),
        (elements[keep], positions[keep]),
    )

def atomic_shape_descriptors(self, l_max=5, radius=6.0, return_coefficients=False) -> np.ndarray:
    """
    Calculate the shape descriptors[1,2] for all symmetry unique
    atoms in this crystal.

    Args:
        l_max (int, optional): maximum level of angular momenta to include in the spherical harmonic
            transform of the molecular shape function.
        radius (float, optional): maximum distance (Angstroms) to include surroundings
            in the shape description
        return_coefficients (bool, optional): also return the spherical harmonic coefficients

    Returns:
        shape description vector

    References:
    ```
    [1] PR Spackman et al. Sci. Rep. 6, 22204 (2016)
        https://dx.doi.org/10.1038/srep22204
    [2] PR Spackman et al. Angew. Chem. 58 (47), 16780-16784 (2019)
        https://dx.doi.org/10.1002/anie.201906602
    ```
    """
    descriptors = []
    coeffs = []
    from chmpy.shape import SHT, stockholder_weight_descriptor

    sph = SHT(l_max=l_max)
    for n, pos, neighbour_els, neighbour_pos in self.atomic_surroundings(
        radius=radius
    ):
        ubound = Element[n].vdw_radius * 3
        desc = stockholder_weight_descriptor(
            sph, [n], [pos], neighbour_els, neighbour_pos, bounds=(0.2, ubound),
            coefficients=return_coefficients
        )
        if return_coefficients:
            descriptors.append(desc[1])
            coeffs.append(desc[0])
        else:
            descriptors.append(desc)
    if return_coefficients:
        return np.asarray(coeffs), np.asarray(descriptors)
    else:
        return np.asarray(descriptors)

def atomic_surroundings(self, radius=6.0) -> List[Tuple]:
    """
    Calculate all atoms within the given `radius` of
    each atomic site in the asymmetric unit.

    Arguments:
        radius (float): the maximum distance (Angstroms) from the origin for inclusion

    Returns:
        A list of atomic number, Cartesian position for both the
        atomic site in question and the surroundings (as an array)
    """
    cart_asym = self.to_cartesian(self.asymmetric_unit.positions)
    hklmax = np.array([-np.inf, -np.inf, -np.inf])
    hklmin = np.array([np.inf, np.inf, np.inf])
    frac_radius = radius / np.array(self.unit_cell.lengths)
    for pos in self.asymmetric_unit.positions:
        hklmax = np.maximum(hklmax, np.ceil(frac_radius + pos))
        hklmin = np.minimum(hklmin, np.floor(pos - frac_radius))
    hmax, kmax, lmax = hklmax.astype(int)
    hmin, kmin, lmin = hklmin.astype(int)
    slab = self.slab(bounds=((hmin, kmin, lmin), (hmax, kmax, lmax)))
    tree = KDTree(slab["cart_pos"])
    results = []
    for n, pos in zip(self.asymmetric_unit.elements, cart_asym):
        idxs = tree.query_ball_point(pos, radius)
        positions = slab["cart_pos"][idxs]
        elements = slab["element"][idxs]
        d = np.linalg.norm(positions - pos, axis=1)
        keep = np.where(d > 1e-3)[0]
        results.append((n.atomic_number, pos, elements[keep], positions[keep]))
    return results

def atoms_in_radius(self, radius, origin=(0, 0, 0)) -> dict:
    """
    Calculate all (periodic) atoms within the given `radius` of the specified
    `origin`.

    Arguments:
        radius (float): the maximum distance (Angstroms) from the origin for inclusion
        origin (Tuple, optional): the origin in fractional coordinates

    Returns:
        A dictionary mapping (see the the `slab` method),
        of those atoms within `radius` of the `origin`.
    """
    frac_origin = self.to_fractional(origin)
    frac_radius = radius / np.array(self.unit_cell.lengths)
    hmax, kmax, lmax = np.ceil(frac_radius + frac_origin).astype(int)
    hmin, kmin, lmin = np.floor(frac_origin - frac_radius).astype(int)
    slab = self.slab(bounds=((hmin, kmin, lmin), (hmax, kmax, lmax)))
    tree = KDTree(slab["cart_pos"])
    idxs = sorted(tree.query_ball_point(origin, radius))
    result = {k: v[idxs] for k, v in slab.items() if isinstance(v, np.ndarray)}
    result["uc_atom"] = np.tile(np.arange(slab["n_uc"]), slab["n_cells"])[idxs]
    return result

def cartesian_symmetry_operations(self):
    """
    Create a list of symmetry operations (rotation, translation)
    for evaluation of transformations in cartesian space.

    The rotation matrices are stored to be used as np.dot(x, R),
    (i.e. post-multiplicaiton on row-major coordinates)

    Returns:
        List[Tuple[np.ndarray, np.ndarray]]: a list of (rotation, translation)
    """
    cart_symops = []
    d = self.unit_cell.direct
    i = self.unit_cell.inverse
    for symop in self.symmetry_operations:
        cart_symops.append(
            (
                np.dot(d.T, np.dot(symop.rotation, i.T)).T,
                self.to_cartesian(symop.translation),
            )
        )
    return cart_symops

def choose_trigonal_lattice(self, choice="H"):
    """
    Change the choice of lattice for this crystal to either
    rhombohedral or hexagonal cell

    Args:
        choice (str, optional): The choice of the resulting lattice, either 'H' for hexagonal
            or 'R' for rhombohedral (default 'H').
    """
    if not self.space_group.has_hexagonal_rhombohedral_choices():
        raise ValueError("Invalid space group for choose_trigonal_lattice")
    if self.space_group.choice == choice:
        return
    cart_asym_pos = self.to_cartesian(self.asymmetric_unit.positions)
    assert choice in ("H", "R"), "Valid choices are H, R"
    if self.space_group.choice == "R":
        T = np.array(((-1, 1, 0), (1, 0, -1), (1, 1, 1)))
    else:
        T = 1 / 3 * np.array(((-1, 1, 1), (2, 1, 1), (-1, -2, 1)))
    new_uc = UnitCell(np.dot(T, self.unit_cell.direct))
    self.unit_cell = new_uc
    self.asymmetric_unit.positions = self.to_fractional(cart_asym_pos)
    self.space_group = SpaceGroup(
        self.space_group.international_tables_number, choice=choice
    )

@classmethod
def from_cif_data(cls, cif_data, titl=None):
    """Initialize a crystal structure from a dictionary
    of CIF data"""
    labels = cif_data.get("atom_site_label", None)
    symbols = cif_data.get("atom_site_type_symbol", None)
    if symbols is None:
        if labels is None:
            raise ValueError(
                "Unable to determine elements in CIF, "
                "need one of _atom_site_label or "
                "_atom_site_type_symbol present"
            )
        elements = [Element[x] for x in labels]
    else:
        elements = [Element[x] for x in symbols]
    x = np.asarray(cif_data.get("atom_site_fract_x", []))
    y = np.asarray(cif_data.get("atom_site_fract_y", []))
    z = np.asarray(cif_data.get("atom_site_fract_z", []))
    occupation = np.asarray(cif_data.get("atom_site_occupancy", [1] * len(x)))
    frac_pos = np.array([x, y, z]).T
    asym = AsymmetricUnit(
        elements=elements, positions=frac_pos, labels=labels, occupation=occupation
    )
    lengths = [cif_data[f"cell_length_{x}"] for x in ("a", "b", "c")]
    angles = [cif_data[f"cell_angle_{x}"] for x in ("alpha", "beta", "gamma")]
    unit_cell = UnitCell.from_lengths_and_angles(lengths, angles, unit="degrees")

    space_group = SpaceGroup(1)
    symop_data_names = (
        "symmetry_equiv_pos_as_xyz",
        "space_group_symop_operation_xyz",
    )
    number = space_group.international_tables_number
    for k in ("space_group_IT_number", "symmetry_Int_Tables_number"):
        if k in cif_data:
            number = cif_data[k]
            break

    for symop_data_block in symop_data_names:
        if symop_data_block in cif_data:
            symops = [
                SymmetryOperation.from_string_code(x)
                for x in cif_data[symop_data_block]
            ]
            try:
                new_sg = SpaceGroup.from_symmetry_operations(symops)
                space_group = new_sg
            except ValueError:
                space_group.symmetry_operations = symops
                symbol = cif_data.get("symmetry_space_group_name_H-M", "Unknown")
                space_group.international_tables_number = number
                space_group.symbol = symbol
                space_group.full_symbol = symbol
                LOG.warn(
                    "Initializing non-standard spacegroup setting %s, "
                    "some SG data may be missing",
                    symbol,
                )
            break
    else:
        # fall back to international tables number
        space_group = SpaceGroup(number)

    return Crystal(unit_cell, space_group, asym, cif_data=cif_data, titl=titl)

@classmethod
def from_cif_file(cls, filename, data_block_name=None):
    """Initialize a crystal structure from a CIF file"""
    cif = Cif.from_file(filename)
    if data_block_name is not None:
        return cls.from_cif_data(cif.data[data_block_name], titl=data_block_name)

    crystals = {
        name: cls.from_cif_data(data, titl=name) for name, data in cif.data.items()
    }
    keys = list(crystals.keys())
    if len(keys) == 1:
        return crystals[keys[0]]
    return crystals

@classmethod
def from_shelx_file(cls, filename, **kwargs):
    """Initialize a crystal structure from a shelx .res file"""
    p = Path(filename)
    titl = p.stem
    return cls.from_shelx_string(p.read_text(), titl=titl, **kwargs)

@classmethod
def from_shelx_string(cls, file_content, **kwargs):
    """Initialize a crystal structure from a shelx .res string"""
    from chmpy.fmt.shelx import parse_shelx_file_content

    shelx_dict = parse_shelx_file_content(file_content)
    asymmetric_unit = AsymmetricUnit.from_records(shelx_dict["ATOM"])
    space_group = SpaceGroup.from_symmetry_operations(
        shelx_dict["SYMM"], expand_latt=shelx_dict["LATT"]
    )
    unit_cell = UnitCell.from_lengths_and_angles(
        shelx_dict["CELL"]["lengths"], shelx_dict["CELL"]["angles"], unit="degrees"
    )
    return cls(unit_cell, space_group, asymmetric_unit, **kwargs)

@classmethod
def from_vasp_file(cls, filename, **kwargs):
    "Initialize a crystal structure from a VASP POSCAR file"
    return cls.from_vasp_string(Path(filename).read_text(), **kwargs)

@classmethod
def from_vasp_string(cls, string, **kwargs):
    "Initialize a crystal structure from a VASP POSCAR string"
    from chmpy.fmt.vasp import parse_poscar

    vasp_data = parse_poscar(string)
    uc = UnitCell(vasp_data["direct"])
    sg = SpaceGroup(1)
    coords = vasp_data["positions"]
    if not vasp_data["coord_type"].startswith("d"):
        coords = uc.to_fractional(coords)
    asym = AsymmetricUnit(vasp_data["elements"], coords)
    return Crystal(uc, sg, asym, titl=vasp_data["name"])

def functional_group_shape_descriptors(
    self, l_max=5, radius=6.0, kind="carboxylic_acid"
) -> np.ndarray:
    """
    Calculate the shape descriptors `[1,2]` for the all atoms in the functional group
    given for all symmetry unique molecules in this crystal.

    Args:
        l_max (int, optional): maximum level of angular momenta to include in the spherical harmonic
            transform of the molecular shape function. (default: 5)
        radius (float, optional): maximum distance (Angstroms) of neighbouring atoms to include in
            stockholder weight calculation (default: 5)
        kind (str, optional): Identifier for the functional group type (default: 'carboxylic_acid')

    Returns:
        shape description vector

    References:
    ```
    [1] PR Spackman et al. Sci. Rep. 6, 22204 (2016)
        https://dx.doi.org/10.1038/srep22204
    [2] PR Spackman et al. Angew. Chem. 58 (47), 16780-16784 (2019)
        https://dx.doi.org/10.1002/anie.201906602
    ```
    """
    descriptors = []
    from chmpy.shape import SHT, stockholder_weight_descriptor

    sph = SHT(l_max=l_max)
    for (
        in_els,
        in_pos,
        neighbour_els,
        neighbour_pos,
    ) in self.functional_group_surroundings(kind=kind, radius=radius):
        masses = np.asarray([Element[x].mass for x in in_els])
        c = np.sum(in_pos * masses[:, np.newaxis] / np.sum(masses), axis=0).astype(
            np.float32
        )
        dists = np.linalg.norm(in_pos - c, axis=1)
        bounds = np.min(dists) / 2, np.max(dists) + 10.0
        descriptors.append(
            stockholder_weight_descriptor(
                sph,
                in_els,
                in_pos,
                neighbour_els,
                neighbour_pos,
                origin=c,
                bounds=bounds,
            )
        )
    return np.asarray(descriptors)

def functional_group_surroundings(self, radius=6.0, kind="carboxylic_acid") -> List:
    """
    Calculate the atomic information for all
    atoms surrounding each functional group in each symmetry unique molecule
    in this crystal within the given radius.

    Args:
        radius (float, optional): Maximum distance in Angstroms between any atom in the molecule
            and the resulting neighbouring atoms
        kind (str, optional): the functional group type

    Returns:
        A list of tuples of (func_el, func_pos, neigh_el, neigh_pos)
        where `func_el` and `neigh_el` are `np.ndarray` of atomic numbers,
        and `func_pos` and `neigh_pos` are `np.ndarray` of Cartesian atomic positions
    """
    results = []
    for mol in self.symmetry_unique_molecules():
        hklmax = np.array([-np.inf, -np.inf, -np.inf])
        hklmin = np.array([np.inf, np.inf, np.inf])
        frac_radius = radius / np.array(self.unit_cell.lengths)
        for pos in self.to_fractional(mol.positions):
            hklmax = np.maximum(hklmax, np.ceil(frac_radius + pos))
            hklmin = np.minimum(hklmin, np.floor(pos - frac_radius))
        hmax, kmax, lmax = hklmax.astype(int)
        hmin, kmin, lmin = hklmin.astype(int)
        slab = self.slab(bounds=((hmin, kmin, lmin), (hmax, kmax, lmax)))
        elements = slab["element"]
        positions = slab["cart_pos"]
        tree = KDTree(positions)
        groups = mol.functional_groups(kind=kind)
        for fg in groups:
            fg = list(fg)
            keep = np.zeros(positions.shape[0], dtype=bool)
            inside = []
            for pos in mol.positions[fg]:
                idxs = tree.query_ball_point(pos, radius)
                d, nn = tree.query(pos)
                keep[idxs] = True
                if d < 1e-3:
                    inside.append(nn)
                    keep[inside] = False
            results.append(
                (
                    mol.atomic_numbers[fg],
                    mol.positions[fg],
                    elements[keep],
                    positions[keep],
                )
            )
    return results

def hirshfeld_surfaces(self, **kwargs):
    "Alias for `self.stockholder_weight_isosurfaces`"
    return self.stockholder_weight_isosurfaces(**kwargs)

@classmethod
def load(cls, filename, **kwargs) -> Union["Crystal", dict]:
    """
    Load a crystal structure from file (.res, .cif)

    Args:
        filename (str): the path to the crystal structure file

    Returns:
        the resulting crystal structure or dictionary of crystal structures
    """
    fpath = Path(filename)
    n = fpath.name
    fname_map = cls._fname_load_map()
    if n in fname_map:
        return fname_map[n](filename)
    extension_map = cls._ext_load_map()
    extension = kwargs.pop("fmt", fpath.suffix.lower())
    if not extension.startswith("."):
        extension = "." + extension
    return extension_map[extension](filename, **kwargs)

def mesh_scene(self, **kwargs):
    """
    Calculate a scene of this meshes of unit cell molecules in this crystal,
    along with optional void surface.

    Args:
        kwargs: optional arguments used in the generation of this scene.

    Returns:
        trimesh.scene.Scene: trimesh scene object.
    """
    from trimesh import Scene

    meshes = {}
    for i, m in enumerate(self.unit_cell_molecules()):
        mesh = m.to_mesh(representation=kwargs.get("representation", "ball_stick"))
        n = m.molecular_formula
        for k, v in mesh.items():
            meshes[f"mol_{i}_{n}.{k}"] = v

    if kwargs.get("void", False):
        void_kwargs = kwargs.get("void_kwargs", {})
        meshes["void_surface"] = self.void_surface(**void_kwargs)
    if kwargs.get("axes", False):
        from trimesh.creation import axis

        meshes["axes"] = axis(
            transform=self.unit_cell.direct_homogeneous.T, axis_length=1.0
        )
    return Scene(meshes)

def molecular_shape_descriptors(
    self, l_max=5, radius=6.0, with_property=None, return_coefficients=False
) -> np.ndarray:
    """
    Calculate the molecular shape descriptors[1,2] for all symmetry unique
    molecules in this crystal.

    Args:
        l_max (int, optional): maximum level of angular momenta to include in the spherical harmonic
            transform of the molecular shape function.
        radius (float, optional): maximum distance (Angstroms) to include surroundings
            in the shape description
        with_property (str, optional): name of the surface property to include in the shape description
        return_coefficients (bool, optional): also return the spherical harmonic coefficients

    Returns:
        shape description vector

    References:
    ```
    [1] PR Spackman et al. Sci. Rep. 6, 22204 (2016)
        https://dx.doi.org/10.1038/srep22204
    [2] PR Spackman et al. Angew. Chem. 58 (47), 16780-16784 (2019)
        https://dx.doi.org/10.1002/anie.201906602
    ```
    """
    descriptors = []
    coeffs = []
    from chmpy.shape import SHT, stockholder_weight_descriptor

    sph = SHT(l_max=l_max)
    for mol, neighbour_els, neighbour_pos in self.molecule_environments(
        radius=radius
    ):
        c = np.array(mol.centroid, dtype=np.float32)
        dists = np.linalg.norm(mol.positions - c, axis=1)
        bounds = np.min(dists) / 2, np.max(dists) + 10.0
        descriptor = stockholder_weight_descriptor(
            sph,
            mol.atomic_numbers,
            mol.positions,
            neighbour_els,
            neighbour_pos,
            origin=c,
            bounds=bounds,
            with_property=with_property,
            coefficients=return_coefficients
        )

        if return_coefficients:
            coeffs.append(descriptor[0])
            descriptors.append(descriptor[1])
        else:
            descriptors.append(descriptor)
    if return_coefficients:
        return np.asarray(coeffs), np.asarray(descriptors)
    else:
        return np.asarray(descriptors)

def molecular_shell(
    self, mol_idx=0, radius=3.8, method="nearest_atom"
) -> List[Molecule]:
    """
    Calculate the neighbouring molecules around the molecule with index
    `mol_idx`, within the given `radius` using the specified `method`.

    Arguments:
        mol_idx (int, optional): The index (into `symmetry_unique_molecules`) of the central
            molecule for the shell
        radius (float, optional): The maximum distance (Angstroms) between the central
            molecule and the neighbours.
        method (str, optional): the method to use when determining inclusion of neighbours.

    Returns:
        A list of neighbouring molecules using the given method.
    """
    mol = self.symmetry_unique_molecules()[mol_idx]
    frac_origin = self.to_fractional(mol.center_of_mass)
    frac_radius = radius / np.array(self.unit_cell.lengths)
    hmax, kmax, lmax = np.ceil(frac_radius + frac_origin).astype(int) + 1
    hmin, kmin, lmin = np.floor(frac_origin - frac_radius).astype(int) - 1
    uc_mols = self.unit_cell_molecules()
    shifts = self.to_cartesian(
        cartesian_product(
            np.arange(hmin, hmax), np.arange(kmin, kmax), np.arange(lmin, lmax)
        )
    )
    neighbours = []
    for uc_mol in uc_mols:
        for shift in shifts:
            uc_mol_t = uc_mol.translated(shift)
            dist = mol.distance_to(uc_mol_t, method=method)
            if (dist < radius) and (dist > 1e-2):
                neighbours.append(uc_mol_t)
    return neighbours

def molecule_dict(self, **kwargs) -> dict:
    """
    A dictionary of `symmetry_unique_molecules`, grouped by
    their chemical formulae.

    Returns:
        the dictionary of molecules with chemical formula keys
        and list of molecule values.
    """
    result = {}
    mols = self.symmetry_unique_molecules()
    for m in mols:
        f = m.molecular_formula
        if f not in result:
            result[f] = []
        result[f].append(m)
    return result

def molecule_environment(self, mol, radius=6.0, threshold=1e-3) -> Tuple:
    """
    Calculate the atomic information for all
    atoms surrounding the given molecule in this crystal
    within the given radius. Atoms closer than `threshold`
    to any atom in the provided molecule will be excluded and
    considered part of the molecule.

    Args:
        mol (Molecule): the molecule whose environment to calculate
        radius (float, optional): Maximum distance in Angstroms between any atom in the molecule
            and the resulting neighbouring atoms
        threshold (float, optional): tolerance for detecting the neighbouring sites as part of the
            given molecule.

    Returns:
        A list of tuples of (Molecule, elements, positions)
            where `elements` is an `np.ndarray` of atomic numbers,
            and `positions` is an `np.ndarray` of Cartesian atomic positions
    """

    hklmax = np.array([-np.inf, -np.inf, -np.inf])
    hklmin = np.array([np.inf, np.inf, np.inf])
    frac_radius = radius / np.array(self.unit_cell.lengths)
    for pos in self.to_fractional(mol.positions):
        hklmax = np.maximum(hklmax, np.ceil(frac_radius + pos))
        hklmin = np.minimum(hklmin, np.floor(pos - frac_radius))
    hmax, kmax, lmax = hklmax.astype(int)
    hmin, kmin, lmin = hklmin.astype(int)
    slab = self.slab(bounds=((hmin, kmin, lmin), (hmax, kmax, lmax)))
    elements = slab["element"]
    positions = slab["cart_pos"]
    tree = KDTree(positions)
    keep = np.zeros(positions.shape[0], dtype=bool)
    this_mol = []
    for pos in mol.positions:
        idxs = tree.query_ball_point(pos, radius)
        d, nn = tree.query(pos)
        keep[idxs] = True
        if d < threshold:
            this_mol.append(nn)
            keep[this_mol] = False
    return (mol, elements[keep], positions[keep])

def molecule_environments(self, radius=6.0, threshold=1e-3) -> List[Tuple]:
    """
    Calculate the atomic information for all
    atoms surrounding each symmetry unique molecule
    in this crystal within the given radius.

    Args:
        radius (float, optional): Maximum distance in Angstroms between any atom in the molecule
            and the resulting neighbouring atoms
        threshold (float, optional): tolerance for detecting the neighbouring sites as part of the
            given molecule.

    Returns:
        A list of tuples of (Molecule, elements, positions)
        where `elements` is an `np.ndarray` of atomic numbers,
        and `positions` is an `np.ndarray` of Cartesian atomic positions
    """
    return [
        self.molecule_environment(x, radius=radius, threshold=threshold)
        for x in self.symmetry_unique_molecules()
    ]

def molecule_shape_descriptors(
    self, mol, l_max=5, radius=6.0, with_property=None
) -> np.ndarray:
    """
    Calculate the molecular shape descriptors `[1,2]` for
    the provided molecule in the crystal.

    Args:
        l_max (int, optional): maximum level of angular momenta to include in the spherical harmonic
            transform of the molecular shape function.
        radius (float, optional): maximum distance (Angstroms) to include surroundings
            in the shape description
        with_property (str, optional): name of the surface property to include in the shape description

    Returns:
        shape description vector

    References:
    ```
    [1] PR Spackman et al. Sci. Rep. 6, 22204 (2016)
        https://dx.doi.org/10.1038/srep22204
    [2] PR Spackman et al. Angew. Chem. 58 (47), 16780-16784 (2019)
        https://dx.doi.org/10.1002/anie.201906602
    ```
    """
    from chmpy.shape import SHT, stockholder_weight_descriptor

    sph = SHT(l_max=l_max)
    mol, neighbour_els, neighbour_pos = self.molecule_environment(
        mol, radius=radius
    )
    c = np.array(mol.centroid, dtype=np.float32)
    dists = np.linalg.norm(mol.positions - c, axis=1)
    bounds = np.min(dists) / 2, np.max(dists) + 10.0
    return stockholder_weight_descriptor(
        sph,
        mol.atomic_numbers,
        mol.positions,
        neighbour_els,
        neighbour_pos,
        origin=c,
        bounds=bounds,
        with_property=with_property,
    )

def promolecule_density_isosurfaces(self, **kwargs) -> List[Trimesh]:
    """
    Calculate promolecule electron density isosurfaces
    for each symmetry unique molecule in this crystal.

    Args:
        kwargs: Keyword arguments used by `Molecule.promolecule_density_isosurface`.

            Options are:
            ```
            isovalue (float, optional): level set value for the isosurface (default=0.002) in au.
            separation (float, optional): separation between density grid used in the surface calculation
                (default 0.2) in Angstroms.
            color (str, optional): surface property to use for vertex coloring, one of ('d_norm_i',
                'd_i', 'd_norm_e', 'd_e')
            colormap (str, optional): matplotlib colormap to use for surface coloring (default 'viridis_r')
            midpoint (float, optional): midpoint of the segmented colormap (if applicable)
            ```

    Returns:
        A list of meshes representing the promolecule density isosurfaces
    """
    return [
        mol.promolecule_density_isosurface(**kwargs)
        for mol in self.symmetry_unique_molecules()
    ]

def save(self, filename, **kwargs):
    """Save this crystal structure to file (.cif, .res, POSCAR)"""
    fpath = Path(filename)
    n = fpath.name
    fname_map = self._fname_save_map()
    if n in fname_map:
        return fname_map[n](filename, **kwargs)
    extension_map = self._ext_save_map()
    extension = kwargs.pop("fmt", fpath.suffix.lower())
    if not extension.startswith("."):
        extension = "." + extension
    return extension_map[extension](filename, **kwargs)

def slab(self, bounds=((-1, -1, -1), (1, 1, 1))) -> dict:
    """
    Calculate the atoms and associated information
    for a slab consisting of multiple unit cells.

    If unit cell atoms have not been calculated, this calculates
    their information and caches it.

    Args:
        bounds (Tuple, optional): Tuple of upper and lower corners (hkl) describing the bounds
            of the slab.

    Returns:
        A dictionary of arrays associated with all sites contained
        in the unit cell of this crystal, members are:
            asym_atom: corresponding asymmetric unit atom indices for all sites.
            frac_pos: (N, 3) array of fractional positions for all sites.
            cart_pos: (N, 3) array of cartesian positions for all sites.
            element: (N) array of atomic numbers for all sites.
            symop: (N) array of indices corresponding to the generator symmetry
                operation for each site.
            label: (N) array of string labels corresponding to each site
            occupation: (N) array of occupation numbers for each site. Will
                warn if any of these are greater than 1.0
            cell: (N,3) array of cell indices for each site

        n_uc: number of atoms in the unit cell

        n_cells: number of cells in this slab

        occupation: (N) array of occupation numbers for each site. Will
        warn if any of these are greater than 1.0

    """
    uc_atoms = self.unit_cell_atoms()
    (hmin, kmin, lmin), (hmax, kmax, lmax) = bounds
    h = np.arange(hmin, hmax + 1)
    k = np.arange(kmin, kmax + 1)
    l = np.arange(lmin, lmax + 1)  # noqa: E741
    cells = cartesian_product(
        h[np.argsort(np.abs(h))], k[np.argsort(np.abs(k))], l[np.argsort(np.abs(l))]
    )
    ncells = len(cells)
    uc_pos = uc_atoms["frac_pos"]
    n_uc = len(uc_pos)
    pos = np.empty((ncells * n_uc, 3), dtype=np.float64)
    slab_cells = np.empty((ncells * n_uc, 3), dtype=np.float64)
    for i, cell in enumerate(cells):
        pos[i * n_uc : (i + 1) * n_uc, :] = uc_pos + cell
        slab_cells[i * n_uc : (i + 1) * n_uc] = cell
    slab_dict = {
        k: np.tile(v, ncells) for k, v in uc_atoms.items() if not k.endswith("pos")
    }
    slab_dict["frac_pos"] = pos
    slab_dict["cell"] = slab_cells
    slab_dict["n_uc"] = n_uc
    slab_dict["n_cells"] = ncells
    slab_dict["cart_pos"] = self.to_cartesian(pos)
    return slab_dict

def stockholder_weight_isosurfaces(self, kind="mol", **kwargs) -> List[Trimesh]:
    """
    Calculate stockholder weight isosurfaces (i.e. Hirshfeld surfaces)
    for each symmetry unique molecule or atom in this crystal.

    Args:
        kind (str, optional): dictates whether we calculate surfaces for each unique molecule
            or for each unique atom
        kwargs: keyword arguments passed to `stockholder_weight_isosurface`.

            Options include:
            ```
            isovalue: float, optional
                level set value for the isosurface (default=0.5). Must be between
                0 and 1, but values other than 0.5 probably won't make sense anyway.
            separation: float, optional
                separation between density grid used in the surface calculation
                (default 0.2) in Angstroms.
            radius: float, optional
                maximum distance for contributing neighbours for the stockholder
                weight calculation
            color: str, optional
                surface property to use for vertex coloring, one of ('d_norm_i',
                'd_i', 'd_norm_e', 'd_e', 'd_norm')
            colormap: str, optional
                matplotlib colormap to use for surface coloring (default 'viridis_r')
            midpoint: float, optional, default 0.0 if using d_norm
                use the midpoint norm (as is used in CrystalExplorer)
            ```

    Returns:
        A list of meshes representing the stockholder weight isosurfaces
    """
    from chmpy import StockholderWeight
    from chmpy.surface import stockholder_weight_isosurface
    from chmpy.util.color import property_to_color
    import trimesh

    sep = kwargs.get("separation", kwargs.get("resolution", 0.2))
    radius = kwargs.get("radius", 12.0)
    vertex_color = kwargs.get("color", "d_norm")
    isovalue = kwargs.get("isovalue", 0.5)
    meshes = []
    extra_props = {}
    isos = []
    if kind == "atom":
        for n, pos, neighbour_els, neighbour_pos in self.atomic_surroundings(
            radius=radius
        ):
            s = StockholderWeight.from_arrays(
                [n], [pos], neighbour_els, neighbour_pos
            )
            iso = stockholder_weight_isosurface(s, isovalue=isovalue, sep=sep)
            isos.append(iso)
    elif kind == "mol":
        for mol, n_e, n_p in self.molecule_environments(radius=radius):
            if vertex_color == "esp":
                extra_props["esp"] = mol.electrostatic_potential
            s = StockholderWeight.from_arrays(
                mol.atomic_numbers, mol.positions, n_e, n_p
            )
            iso = stockholder_weight_isosurface(
                s, isovalue=isovalue, sep=sep, extra_props=extra_props
            )
            isos.append(iso)
    else:
        for arr in self.functional_group_surroundings(radius=radius, kind=kind):
            s = StockholderWeight.from_arrays(*arr)
            iso = stockholder_weight_isosurface(s, isovalue=isovalue, sep=sep)
            isos.append(iso)

    for iso in isos:
        prop = iso.vertex_prop[vertex_color]
        color = property_to_color(prop, cmap=kwargs.get("cmap", vertex_color))
        mesh = trimesh.Trimesh(
            vertices=iso.vertices,
            faces=iso.faces,
            normals=iso.normals,
            vertex_colors=color,
        )
        meshes.append(mesh)
    return meshes

def symmetry_unique_dimers(self, radius=3.8, distance_method="nearest_atom"):
    """
    Calculate the information for all molecule
    pairs surrounding the symmetry_unique_molecules
    in this crystal within the given radius. 

    Args:
        radius (float, optional): Maximum distance in Angstroms between any atom in the molecule
            and the resulting neighbouring atoms

    Returns:
        A dictionary of dimers (Molecule, elements, positions)
            where `elements` is an `np.ndarray` of atomic numbers,
            and `positions` is an `np.ndarray` of Cartesian atomic positions
    """
    from chmpy.core.dimer import Dimer
    from copy import deepcopy
    from collections import defaultdict

    hklmax = np.array([-np.inf, -np.inf, -np.inf])
    hklmin = np.array([np.inf, np.inf, np.inf])
    frac_radius = radius / np.array(self.unit_cell.lengths)

    for pos in self.asymmetric_unit.positions:
        hklmax = np.maximum(hklmax, np.ceil(frac_radius + pos))
        hklmin = np.minimum(hklmin, np.floor(pos - frac_radius))

    hmax, kmax, lmax = hklmax.astype(int)
    hmin, kmin, lmin = hklmin.astype(int)

    shifts_frac = cartesian_product(
        np.arange(hmin, hmax), np.arange(kmin, kmax), np.arange(lmin, lmax)
    )

    shifts = self.to_cartesian(shifts_frac)
    LOG.debug(
        "Looking in %d neighbouring cells: %s : %s",
        len(shifts),
        hklmin.astype(int),
        hklmax.astype(int),
    )
    all_dimers = []
    mol_dimers = []
    for mol_a in self.symmetry_unique_molecules():
        dimers_a = []
        for mol_b in self.unit_cell_molecules():
            for shift, shift_frac in zip(shifts, shifts_frac):
                # shift_frac assumes the molecule is generated from the [0, 0, 0] cell, it's not
                mol_bt = mol_b.translated(shift)
                r = mol_a.distance_to(mol_bt, method=distance_method)
                if r > 1e-1 and r < radius:
                    d = Dimer(
                        mol_a,
                        mol_bt,
                        separation=r,
                        transform_ab="calculate",
                        frac_shift=shift_frac,
                    )
                    for i, dimer in enumerate(all_dimers):
                        if dimer.separation <= d.separation + 1e-3:
                            if d == dimer:
                                dimers_a.append(i)
                                break
                    else:
                        dimers_a.append(len(all_dimers))
                        all_dimers.append(d)
                    all_dimers = sorted(all_dimers, key=lambda x: x.separation)
        mol_dimers.append(dimers_a)
    return all_dimers, mol_dimers

def symmetry_unique_molecules(self, bond_tolerance=0.4) -> List[Molecule]:
    """
    Calculate a list of connected molecules which contain
    every site in the asymmetric_unit

    Populates the _symmetry_unique_molecules member, subsequent
    calls to this function will be a no-op.

    Args:
        bond_tolerance (float, optional): Bonding tolerance (bonded if d < cov_a + cov_b + bond_tolerance)

    Returns:
        List of all connected molecules in the asymmetric_unit of this
        crystal, i.e. the minimum list of connected molecules which contain
        all sites in the asymmetric unit.

        If the asymmetric is molecular, the list will be of length
        num_molecules_in_asymmetric_unit and the total number of atoms
        will be equal to the number of atoms in the asymmetric_unit
    """

    if hasattr(self, "_symmetry_unique_molecules"):
        return getattr(self, "_symmetry_unique_molecules")
    uc_molecules = self.unit_cell_molecules()
    asym_atoms = np.zeros(len(self.asymmetric_unit), dtype=bool)
    molecules = []

    # sort by % of identity symop
    def order(x):
        return len(np.where(x.asym_symops == 16484)[0]) / len(x)

    for i, mol in enumerate(sorted(uc_molecules, key=order, reverse=True)):
        asym_atoms_in_g = np.unique(mol.properties["asymmetric_unit_atoms"])
        if np.all(asym_atoms[asym_atoms_in_g]):
            continue
        asym_atoms[asym_atoms_in_g] = True
        molecules.append(mol)
        if np.all(asym_atoms):
            break
    LOG.debug("%d symmetry unique molecules", len(molecules))
    setattr(self, "_symmetry_unique_molecules", molecules)
    for i, mol in enumerate(molecules):
        mol.properties["asym_mol_idx"] = i

    ak = "asymmetric_unit_atoms"
    for mol in self.unit_cell_molecules():
        if "asym_mol_idx" in mol.properties:
            continue
        else:
            for asym_mol in molecules:
                if np.all(mol.properties[ak] == asym_mol.properties[ak]):
                    mol.properties["asym_mol_idx"] = asym_mol.properties[
                        "asym_mol_idx"
                    ]
                    break
            else:
                LOG.warn(
                    "No equivalent asymmetric unit molecule found!? -- this should not happen!"
                )
    return molecules

def to_cartesian(self, coords) -> np.ndarray:
    """
    Convert coordinates (row major) from fractional to cartesian coordinates.

    Arguments:
        coords (np.ndarray): (N, 3) array of positions assumed to be in fractional coordinates

    Returns:
        (N, 3) array of positions transformed to cartesian (orthogonal) coordinates
        by the unit cell of this crystal.
    """
    return self.unit_cell.to_cartesian(coords)