minor refactorings,

making the diversity network a bit more flexible and robust.

src/konnektor/network_planners/concatenator/mst_concatenator.py

@@ -1,6 +1,6 @@
 import logging
 import itertools
-from typing import List
+from typing import Iterable
 from konnektor.utils import LigandNetwork
 
 from ..generators.netx_netgen import MstNetworkGenerator
@@ -13,7 +13,7 @@ def __init__(self, mapper, scorer,  n_connecting_edges: int = 3):
         self.scorer = scorer
         self.n_connecting_edges = n_connecting_edges
         self.network_generator = MstNetworkGenerator()
-    def concatenate_networks(self, ligand_networks: List[LigandNetwork]) -> LigandNetwork:
+    def concatenate_networks(self, ligand_networks: Iterable[LigandNetwork]) -> LigandNetwork:
         """
         TODO Separate networking from Ligand stuff
         Parameters

src/konnektor/network_planners/generators/diversity_network_planner.py

@@ -1,58 +1,79 @@
+import logging
+import inspect
 from typing import Iterable, Callable, Union
 
-from gufe import SmallMoleculeComponent, LigandNetwork
-from ._abstract_ligand_network_planner import LigandNetworkPlanner
+from tqdm import tqdm
 
+# Clustering
 from sklearn.cluster import KMeans
 from scikit_mol.fingerprints import RDKitFingerprintTransformer, MorganFingerprintTransformer
 from scikit_mol.descriptors import MolecularDescriptorTransformer
 
-from ...network_tools.cluster_molecules import CompoundDiversityClustering
+from gufe import SmallMoleculeComponent, LigandNetwork
+from ...network_tools import cluster_compound, append_node, concatenate_networks
+from .. import MinimalSpanningTreeLigandNetworkPlanner, CyclicLigandNetworkPlanner
 from ..concatenator import MstConcatenate
+from ._abstract_ligand_network_planner import LigandNetworkPlanner
+
+log = logging.getLogger()
+
 
 class DiversityNetworkPlanner(LigandNetworkPlanner):
     def __init__(self, mapper, scorer,
                  node_featurizer=RDKitFingerprintTransformer(),
-                 clustering= KMeans(n_clusters=3)):
+                 clustering=KMeans(n_clusters=3),
+                 sub_network_planners: Iterable[LigandNetworkPlanner] = (
+                         CyclicLigandNetworkPlanner, MinimalSpanningTreeLigandNetworkPlanner),
+                 concatenator: MstConcatenate = MstConcatenate,
+                 ):
         super().__init__(mapper=mapper, scorer=scorer,
-                       network_generator=None)
-        self.feat =  node_featurizer
-        self.cluster =clustering
+                         network_generator=None)
+        self.feat = node_featurizer
+        self.cluster = clustering
+
+        self.sub_network_planners = [c(mapper=mapper, scorer=scorer) if inspect.isclass(c) else c
+                                     for c in sub_network_planners]
+        self.concatenator = concatenator(mapper=mapper, scorer=scorer) if inspect.isclass(
+            concatenator) else concatenator
 
     def generate_ligand_network(self,
-                         nodes: Iterable[SmallMoleculeComponent],
-                         progress: Union[bool, Callable[[Iterable], Iterable]] = True,
-                         # allow_disconnected=True
-                         ):
-        from konnektor.network_planners import MinimalSpanningTreeLigandNetworkPlanner, CyclicLigandNetworkPlanner
-
-        #step 1: Seperate nodes by diversity
-        cc = CompoundDiversityClustering(featurize=self.feat,
+                                nodes: Iterable[SmallMoleculeComponent],
+                                progress: Union[bool, Callable[[Iterable], Iterable]] = True,
+                                # allow_disconnected=True
+                                ):
+
+        # Step 1: Seperate nodes by diversity
+        self.clusters = cluster_compound(compounds=nodes,
+                                         featurize=self.feat,
                                          cluster=self.cluster)
 
-        clusters = cc.cluster_compounds(nodes)
-
-
-        # Sub Network, based on clusters
-        planner = CyclicLigandNetworkPlanner(mapper=self.mapper, scorer=self.scorer)
-        alt_planner = MinimalSpanningTreeLigandNetworkPlanner(mapper=self.mapper, scorer=self.scorer)
-
-        sub_networks = []
-        for cID, mols in clusters.items():
-            if(len(mols)>1):
-                if(len(mols)>2):
-                    sub_network = planner(mols)
-                    sub_networks.append(sub_network)
-                else:
-                    sub_network = alt_planner(mols)
-                    sub_networks.append(sub_network)
-            else:# Need to generate the Empty Network here!
-                continue
-                sub_network = LigandNetwork(edges=set([]), nodes=mols)
-                sub_networks.append_node(sub_network)
-
-        # Connect the Networks:
-        con = MstConcatenate(mapper=self.mapper, scorer=self.scorer, n_connecting_edges=3)
-        concat_network = con.concatenate_networks(ligand_networks=sub_networks)
+        # Step 2:  Sub Network, based on clusters
+
+        self.sub_networks = []
+        for cID, mols in tqdm(self.clusters.items(), desc="Build Cluster Networks"):
+            if (cID >= 0):  # Noise cluster is not for subnetworks
+                if (len(mols) > 1):
+                    for network_planner in self.sub_network_planners:
+                        try:
+                            sub_network = network_planner(mols)
+                            self.sub_networks.append(sub_network)
+                            break
+                        except Exception as err:
+                            print("ERR", "\n".join(err.args))
+                            continue
+                else:  # Need to generate the Empty Network here!
+                    sub_network = LigandNetwork(edges=set([]), nodes=mols)
+                    self.sub_networks.append(sub_network)
+
+        # step 3: Connect the Networks:
+        log.info("Concatenate Networks")
+        concat_network = concatenate_networks(networks=self.sub_networks,
+                                              concatenator=self.concatenator)
+
+        # step 4: has the clustering a noise cluster
+        if -1 in self.clusters:
+            for mol in tqdm(self.clusters[-1], desc="add Noise Mols"):
+                concat_network = append_node(network=concat_network, compound=mol,
+                                             concatenator=self.concatenator)
 
         return concat_network

src/konnektor/network_tools/__init__.py

@@ -1,3 +1,5 @@
-from .concatenate import concatenate, append_node
+from .concatenate import concatenate_networks, append_node
 from .merge import merge_two_networks
 from .misc import delete_transformation, cluster_compound
+
+

src/konnektor/network_tools/concatenate.py

@@ -1,10 +1,11 @@
+from typing import Iterable
 from gufe import LigandNetwork, SmallMoleculeComponent
 
 from ..network_planners.concatenator.mst_concatenator import MstConcatenate
 
-def concatenate(network1:LigandNetwork,
-                    network2:LigandNetwork,
-           concatenator:MstConcatenate=MstConcatenate)->LigandNetwork:
+
+def concatenate_networks(networks:Iterable[LigandNetwork],
+                         concatenator)->LigandNetwork:
     """
     Merging networks, is similar to a union of a set of nodes and edgees,
     if they are all connected via at least one edge. This means, that  at
@@ -23,14 +24,14 @@ def concatenate(network1:LigandNetwork,
         returns the concatenated network
     """
 
-    concatenated_network = concatenator.concatenate_networks([network1,
-                                                             network2])
+    concat_network = concatenator.concatenate_networks(ligand_networks=networks)
+
+    return concat_network
 
-    return concatenated_network
 
-def append_node(network :LigandNetwork,
-                compound :SmallMoleculeComponent,
-                concatenator :MstConcatenate =MstConcatenate)->LigandNetwork:
+def append_node(network: LigandNetwork,
+                compound: SmallMoleculeComponent,
+                concatenator: MstConcatenate)->LigandNetwork:
     """
     Merging networks, is similar to a union of a set of nodes and edgees,
     if they are all connected via at least one edge. This means, that  at

src/konnektor/network_tools/misc.py

@@ -1,3 +1,5 @@
+from typing import Iterable
+
 from sklearn.base import TransformerMixin, ClusterMixin
 from sklearn.cluster import KMeans
 from scikit_mol.fingerprints import RDKitFingerprintTransformer, MorganFingerprintTransformer
@@ -19,11 +21,12 @@ def delete_transformation(network :LigandNetwork,
     return LigandNetwork(edges=filtered_edges, nodes=network.nodes)
 
 
-def cluster_compound(compounds: list[SmallMoleculeComponent],
+def cluster_compound(compounds: Iterable[SmallMoleculeComponent],
                      featurize:TransformerMixin = RDKitFingerprintTransformer(),
                     cluster:ClusterMixin = KMeans(n_clusters=5, n_init="auto")
                      )->dict[int,
 list[SmallMoleculeComponent]]:
+
     clusterer = CompoundDiversityClustering( featurize = featurize,cluster=cluster)
     return clusterer.cluster_compounds(compounds)