prepare new version

src/moffragmentor/fragmentor/__init__.py

@@ -1,24 +1,36 @@
 # -*- coding: utf-8 -*-
 """Methods for the fragmentation of MOFs"""
-from collections import defaultdict, namedtuple
+from collections import namedtuple
 
 from loguru import logger
 from skspatial.objects import Points
 
 from moffragmentor.sbu.linkercollection import LinkerCollection
 
-from ._no_core_linker import generate_new_node_collection
-from .linkerlocator import create_linker_collection
-from .nodelocator import create_node_collection, find_node_clusters
-from .solventlocator import get_all_bound_solvent_molecules, get_floating_solvent_molecules
-from ..net import build_net
-from ..utils import _get_metal_sublist
+from moffragmentor.fragmentor._no_core_linker import generate_new_node_collection
+from moffragmentor.fragmentor.linkerlocator import create_linker_collection
+from moffragmentor.fragmentor.nodelocator import create_node_collection, find_node_clusters
+from moffragmentor.fragmentor.solventlocator import (
+    get_all_bound_solvent_molecules,
+    get_floating_solvent_molecules,
+)
+from moffragmentor.net import build_net
+from moffragmentor.utils import _get_metal_sublist
+from moffragmentor.descriptors.sbu_dimensionality import get_sbu_dimensionality
 
 __all__ = ["FragmentationResult"]
 
 FragmentationResult = namedtuple(
     "FragmentationResult",
-    ["nodes", "linkers", "bound_solvent", "unbound_solvent", "capping_molecules", "net_embedding"],
+    [
+        "nodes",
+        "linkers",
+        "bound_solvent",
+        "unbound_solvent",
+        "capping_molecules",
+        "net_embedding",
+        "has_1d_sbu",
+    ],
 )
 
 
@@ -29,33 +41,44 @@ def metal_and_branching_coplanar(node, mof, tol=0.1):
     return points.are_coplanar(tol=tol)
 
 
-def run_fragmentation(mof) -> FragmentationResult:  # pylint: disable=too-many-locals
+def run_fragmentation(
+    mof, check_dimensionality: bool = False, create_single_metal_bus: bool = False
+) -> FragmentationResult:
     """Take a MOF and split it into building blocks."""
     logger.debug("Starting fragmentation with location of unbound solvent")
     unbound_solvent = get_floating_solvent_molecules(mof)
     need_rerun = True
     forbidden_indices = []
     counter = 0
-    while need_rerun: 
+    has_1d_sbu = None
+    while need_rerun:
         try:
             not_node = []
             logger.debug(f"Fragmenting MOF for the {counter} time")
             # Find nodes
             node_result = find_node_clusters(
                 mof, unbound_solvent.indices, forbidden_indices=forbidden_indices
             )
+
+            if create_single_metal_bus:
+                # Rewrite the node result
+                ...
+            if check_dimensionality:
+                ...
+                # If we have nodes with dimensionality >0, we change these nodes to only contain the metal
+
             node_collection = create_node_collection(mof, node_result)
+
             # Find bound solvent
             logger.debug("Locating bound solvent")
             bound_solvent = get_all_bound_solvent_molecules(mof, node_result.nodes)
             logger.debug(f"Found bound solvent {len(bound_solvent.indices)>0}")
             # Filter the linkers (valid linkers have at least two branch points)
-        
+
             logger.debug("Locating linkers")
             linker_collection = create_linker_collection(
                 mof, node_result, node_collection, unbound_solvent, bound_solvent
             )
-
 
             logger.debug("Checking for metal in linker")
             # ToDo: factor this out into its own function
@@ -71,7 +94,9 @@ def run_fragmentation(mof) -> FragmentationResult:  # pylint: disable=too-many-l
                         node_collection[i]._original_indices,
                         node_collection[i]._original_graph_branching_indices,
                     )
-                    if metal_and_branching_coplanar(node_collection[i], mof):
+                    if metal_and_branching_coplanar(node_collection[i], mof) & (
+                        len(mof.get_neighbor_indices(metal_in_node[0])) > 2
+                    ):
                         logger.debug(
                             "Metal in linker found, current node: {}, indices: {}".format(
                                 node,
@@ -98,7 +123,7 @@ def run_fragmentation(mof) -> FragmentationResult:  # pylint: disable=too-many-l
             counter += 1
         except Exception as e:
             logger.exception(f"Error while fragmenting: {e}")
-    
+
     logger.debug(
         "Check if we need to move the capping molecules from the linkercollection into their own collection"
     )
@@ -176,6 +201,7 @@ def run_fragmentation(mof) -> FragmentationResult:  # pylint: disable=too-many-l
         unbound_solvent,
         capping_molecules,
         net_embedding,
+        has_1d_sbu,
     )
 
     return fragmentation_results

src/moffragmentor/fragmentor/nodelocator.py

@@ -12,14 +12,14 @@
 import networkx as nx
 from loguru import logger
 
-from ._graphsearch import (
+from moffragmentor.fragmentor._graphsearch import (
     _complete_graph,
     _to_graph,
     recursive_dfs_until_branch,
     recursive_dfs_until_cn3,
 )
-from ..sbu import Node, NodeCollection
-from ..utils import _flatten_list_of_sets
+from moffragmentor.sbu import Node, NodeCollection
+from moffragmentor.utils import _flatten_list_of_sets
 
 __all__ = [
     "find_node_clusters",
@@ -66,7 +66,7 @@ def _path_without_metal_and_branching_sites(path, metal_indices, branching_indic
     return [i for i in path if i not in metal_indices and i not in branching_indices]
 
 
-def find_node_clusters(  # pylint:disable=too-many-locals
+def find_node_clusters(
     mof,
     unbound_solvent_indices: Optional[List[int]] = None,
     forbidden_indices: Optional[List[int]] = None,

src/moffragmentor/sbu/linker.py

@@ -11,6 +11,9 @@
 class Linker(SBU):
     """Describe a linker in a MOF"""
 
+    def __repr__(self) -> str:
+        return f"Linker ({self.composition})"
+
 
 def _get_edge_dict_from_rdkit_mol(mol):
     edges = {}

src/moffragmentor/sbu/node.py

@@ -138,3 +138,6 @@ def from_mof_and_indices(  # pylint:disable=too-many-arguments
             binding_indices,
             connecting_paths=connecting_paths,
         )
+
+    def __repr__(self) -> str:
+        return f"Node ({self.composition})"

src/moffragmentor/sbu/sbucollection.py

@@ -8,6 +8,7 @@
 from .sbu import SBU
 from ..utils import get_molecule_mass
 
+
 class SBUCollection:
     """Container for a collection of SBUs"""
 

src/moffragmentor/utils/__init__.py

@@ -16,12 +16,14 @@
 from pymatgen.core import Molecule, Structure
 from pymatgen.io.babel import BabelMolAdaptor
 
+
 def get_molecule_mass(molecule):
     mass = 0
     for site in molecule:
         mass += site.specie.atomic_mass
     return mass
 
+
 def unwrap(positions, lattice):
     celldiag = np.diagonal(lattice.matrix)
     dxyz = positions - positions[0]
@@ -250,6 +252,7 @@ def _not_relevant_structure_indices(structure: Structure, indices: Collection[in
             not_relevant.append(i)
     return not_relevant
 
+
 def get_sub_structure(mof: "MOF", indices: Collection[int]) -> Structure:
     """Return a sub-structure of the structure with only the sites with the given indices.
 
@@ -267,6 +270,7 @@ def get_sub_structure(mof: "MOF", indices: Collection[int]) -> Structure:
     s = Structure.from_sites(sites, to_unit_cell=True)
     return s
 
+
 def visualize_part(mof, indices: Collection):
     import nglview  # pylint:disable=import-outside-toplevel
 

tests/conftest.py

@@ -3,14 +3,14 @@
 import os
 
 import networkx as nx
+import numpy as np
 import pytest
 from pymatgen.analysis.graphs import MoleculeGraph, StructureGraph
 from pymatgen.analysis.local_env import JmolNN
 from pymatgen.core import Molecule, Structure
 
 from moffragmentor import MOF
 from moffragmentor.sbu import Linker, Node
-import numpy as np
 
 THIS_DIR = os.path.dirname(os.path.realpath(__file__))
 

tests/test_sbu_collection.py

@@ -7,4 +7,4 @@ def test_sbu_collection(get_linker_object):
     assert sbu_collection.sbus[0].composition == "C9 H3 O6"
 
     sbu_collection = SBUCollection([get_linker_object, get_linker_object])
-    assert len(sbu_collection.sbus) == 2
+    assert len(sbu_collection.sbus) == 2

tests/utils/test_utils.py

@@ -1,8 +1,10 @@
-from moffragmentor.utils import get_molecule_mass
-from pymatgen.core import Molecule
 import pytest
+from pymatgen.core import Molecule
 
-def test_get_molecule_mass(get_methane_molecule_and_graph): 
+from moffragmentor.utils import get_molecule_mass
+
+
+def test_get_molecule_mass(get_methane_molecule_and_graph):
     mol, graph = get_methane_molecule_and_graph
-   
-    assert get_molecule_mass(mol) == pytest.approx(16.04303, 0.1)
+
+    assert get_molecule_mass(mol) == pytest.approx(16.04303, 0.1)