molcraft 0.1.0a3__py3-none-any.whl → 0.1.0a5__py3-none-any.whl

molcraft/models.py

@@ -270,6 +270,19 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
         """
         super().load_weights(filepath, *args, **kwargs)
 
+    def embedding(self) -> 'FunctionalGraphModel':
+        model = self
+        if not isinstance(model, FunctionalGraphModel):
+            raise ValueError(
+                'Currently, to extract the embedding part of the model, '
+                'it needs to be a `FunctionalGraphModel`. '
+            )
+        inputs = model.input 
+        for layer in model.layers:
+            if isinstance(layer, layers.Readout):
+                outputs = layer.output 
+        return self.__class__(inputs, outputs, name=f'{self.name}_embedding')
+    
     def train_step(self, tensor: tensors.GraphTensor) -> dict[str, float]:
         y = tensor.context.get('label')
         sample_weight = tensor.context.get('weight')
@@ -381,7 +394,7 @@ def interpret(
 
 def saliency(
     model: GraphModel,
-    graph_tensor: tensors.GraphTensor
+    graph_tensor: tensors.GraphTensor,
 ) -> tensors.GraphTensor:
     x = graph_tensor
     if tensors.is_ragged(x):

molcraft/ops.py

@@ -19,9 +19,16 @@ def gather(
 def aggregate(
     node_feature: tf.Tensor, 
     edge: tf.Tensor, 
-    num_nodes: tf.Tensor
+    num_nodes: tf.Tensor,
+    mode: str = 'sum',
 ) -> tf.Tensor:
-    return keras.ops.segment_sum(node_feature, edge, num_nodes)
+    if mode == 'mean':
+        return segment_mean(
+            node_feature, edge, num_nodes, sorted=False
+        )
+    return keras.ops.segment_sum(
+        node_feature, edge, num_nodes, sorted=False
+    )
 
 def propagate(
     node_feature: tf.Tensor,
@@ -82,7 +89,11 @@ def segment_mean(
     sorted: bool = False,
 ) -> tf.Tensor:
     if num_segments is None:
-        num_segments = keras.ops.max(segment_ids) + 1
+        num_segments = keras.ops.cond(
+            keras.ops.shape(segment_ids)[0] > 0,
+            lambda: keras.ops.max(segment_ids) + 1,
+            lambda: 0
+        )
     if backend.backend() == 'tensorflow':
         return tf.math.unsorted_segment_mean(
             data=data,

molcraft/tensors.py

@@ -219,13 +219,13 @@ class GraphTensor(tf.experimental.BatchableExtensionType):
             raise ValueError
         return ops.gather(self.node[node_attr], self.edge[edge_type])
     
-    def aggregate(self, edge_attr: str, edge_type: str = 'target') -> tf.Tensor:
+    def aggregate(self, edge_attr: str, edge_type: str = 'target', mode: str = 'sum') -> tf.Tensor:
         if edge_type != 'source' and edge_type != 'target':
-            raise ValueError
+            raise ValueError('`edge_attr` needs to be `source` or `target`.')
         edge_attr = self.edge[edge_attr]
         if 'weight' in self.edge:
             edge_attr = edge_attr * self.edge['weight']
-        return ops.aggregate(edge_attr, self.edge[edge_type], self.num_nodes)
+        return ops.aggregate(edge_attr, self.edge[edge_type], self.num_nodes, mode=mode)
     
     def propagate(self, add_edge_feature: bool = False):
         updated_feature = ops.propagate(

{molcraft-0.1.0a3.dist-info → molcraft-0.1.0a5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a3
+Version: 0.1.0a5
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -51,11 +51,11 @@ Dynamic: license-file
 
 ## Highlights
 - Compatible with **Keras 3**
-- Simplified API
-- Fast featurization
-- Modular graph **layers**
-- Serializable graph **featurizers** and **models**
-- Flexible **GraphTensor**
+- Customizable and serializable **featurizers**
+- Customizable and serializable **layers** and **models**
+- Customizable **GraphTensor**
+- Fast and efficient featurization of molecular graphs
+- Efficient and easy-to-use input pipelines using TF **records**
 
 ## Examples 
 

molcraft-0.1.0a5.dist-info/RECORD

@@ -0,0 +1,18 @@
+molcraft/__init__.py,sha256=eTGjgMlXf3I8ThkUwgdiONb5Yc-5fWOFvY8U8WXOMwc,435
+molcraft/callbacks.py,sha256=mkz4ALjJFPy8nHd2nCAuMbKceKnq4tIpZhUuUOvie2Y,1209
+molcraft/chem.py,sha256=_UO5O-I7KUtGf3vRrFEYoAUGlW5xi2x8ylu5f-Ybumo,18696
+molcraft/conformers.py,sha256=p09gOQOdxLSj3yohZOMkxxLriHsZ1ZqOoiWLi73OpIg,4325
+molcraft/datasets.py,sha256=rFgXTC1ZheLhfgQgcCspP_wEE54a33PIneH7OplbS-8,4047
+molcraft/descriptors.py,sha256=gKqlJ3BqJLTeR2ft8isftSEaJDC8cv64eTq5IYhy4XM,3032
+molcraft/features.py,sha256=69oV_GHNdBKPA4sp6Tpo6brvNmaauk_IVIzNjX7VDmg,13648
+molcraft/featurizers.py,sha256=kV5RN_Z2pELjDcwE65KYy_JagbDUueXoClpsIOFsI9I,27073
+molcraft/layers.py,sha256=y-sBLXWttr-fkGZ-acL1srMB8QqeXnHotYK9KCcyJNU,70581
+molcraft/models.py,sha256=0MN4PAlsacni7RfIcYm_imxuzBVL2K8w3MnaUM24DeI,18021
+molcraft/ops.py,sha256=uSnBYQwxYJ1ATdDpr290bxiyQZkrSCVxlB7btlh_n2I,4112
+molcraft/records.py,sha256=w4-bcWZEC0oVInrE1e0kQBroIaSCA0PN1JBPOtO6VUY,5251
+molcraft/tensors.py,sha256=8hwlad000wQ5pNLSdzd3rCXVbaUHBxUq2MbBx27dKzU,22391
+molcraft-0.1.0a5.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a5.dist-info/METADATA,sha256=mb5KnvJUzofmx-MNraJxyiBBug2QNIQTQDGyC1L3SDw,4201
+molcraft-0.1.0a5.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+molcraft-0.1.0a5.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a5.dist-info/RECORD,,

{molcraft-0.1.0a3.dist-info → molcraft-0.1.0a5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

molcraft/experimental/__init__.py

@@ -1 +0,0 @@
-from molcraft.experimental import peptides

molcraft/experimental/peptides.py

@@ -1,264 +0,0 @@
-import re
-import keras
-import numpy as np
-import tensorflow as tf
-from rdkit import Chem
-
-from molcraft import ops
-from molcraft import chem 
-from molcraft import features 
-from molcraft import featurizers
-from molcraft import tensors
-from molcraft import descriptors
-
-    
-def Graph(
-    inputs,
-    atom_features: list[features.Feature] | str | None = 'auto',
-    bond_features: list[features.Feature] | str | None = 'auto',
-    super_atom: bool = True,
-    radius: int | float | None = None,
-    self_loops: bool = False,
-    include_hs: bool = False,
-    **kwargs,
-):
-    featurizer = featurizers.MolGraphFeaturizer(
-        atom_features=atom_features,
-        bond_features=bond_features,
-        molecule_features=[AminoAcidType()],
-        super_atom=super_atom,
-        radius=radius,
-        self_loops=self_loops,
-        include_hs=include_hs,
-        **kwargs,
-    )
-
-    inputs = [
-        residues[x] for x in ['G'] + inputs
-    ]
-    tensor_list = [featurizer(x) for x in inputs]
-    return tf.stack(tensor_list, axis=0)
-
-
-def GraphLookup(graph: tensors.GraphTensor) -> 'GraphLookupLayer':
-    lookup = GraphLookupLayer()
-    lookup._build(graph)
-    return lookup
-
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class GraphLookupLayer(keras.layers.Layer):
-
-    def call(self, indices: tf.Tensor) -> tensors.GraphTensor:
-        indices = tf.sort(tf.unique(tf.reshape(indices, [-1]))[0])
-        graph = self.graph[indices]
-        sizes = graph.context['size']
-        max_index = keras.ops.max(indices)
-        sizes = tf.tensor_scatter_nd_update(
-            tensor=tf.zeros([max_index + 1], dtype=indices.dtype),
-            indices=indices[:, None],
-            updates=sizes
-        )
-        graph = graph.update(
-            {
-                'context': {
-                    'size': sizes
-                }
-            },
-        )
-        return tensors.to_dict(graph)
-    
-    def _build(self, x):
-
-        if isinstance(x, tensors.GraphTensor):
-            tensor = tensors.to_dict(x)
-            self._spec = tf.nest.map_structure(
-                tf.type_spec_from_value, tensor
-            )
-        else:
-            self._spec = x
-
-        self._graph = tf.nest.map_structure(
-            lambda s: self.add_weight(
-                shape=s.shape, 
-                dtype=s.dtype, 
-                trainable=False,
-                initializer='zeros'
-            ),
-            self._spec
-        )
-
-        if isinstance(x, tensors.GraphTensor):
-            tf.nest.map_structure(
-                lambda v, x: v.assign(x),
-                self._graph, tensor
-            )
-
-        graph = tf.nest.map_structure(
-            keras.ops.convert_to_tensor, self._graph
-        )
-        self._graph_tensor = tensors.from_dict(graph)
-        
-    def get_config(self):
-        config = super().get_config()
-        spec = keras.saving.serialize_keras_object(self._spec)
-        config['spec'] = spec
-        return config 
-    
-    @classmethod
-    def from_config(cls, config: dict) -> 'GraphLookupLayer':
-        spec = config.pop('spec')
-        spec = keras.saving.deserialize_keras_object(spec)
-        layer = cls(**config)
-        layer._build(spec)
-        return layer
-
-    @property 
-    def graph(self) -> tensors.GraphTensor:
-        return self._graph_tensor
-    
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class Gather(keras.layers.Layer):
-
-    def __init__(
-        self, 
-        padding: list[tuple[int]] | tuple[int] | int = 1, 
-        mask_value: int = 0,
-        **kwargs
-    ) -> None:
-        super().__init__(**kwargs)
-        self.padding = padding
-        self.mask_value = mask_value
-        self.supports_masking = True
-
-    def get_config(self):
-        config = super().get_config()
-        config['mask_value'] = self.mask_value 
-        config['padding'] = self.padding
-        return config 
-    
-    def call(self, inputs: tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
-        data, indices = inputs 
-        # if self.padding:
-        #     padding = self.padding
-        #     if isinstance(self.padding, int):
-        #         padding = [(self.padding, 0)]
-        #     if isinstance(self.padding, tuple):
-        #         padding = [self.padding] 
-        #     data_rank = len(keras.ops.shape(data))
-        #     for _ in range(data_rank - len(padding)):
-        #         padding.append((0, 0))
-        #     data = keras.ops.pad(data, padding)
-        return ops.gather(data, indices) 
-    
-    def compute_mask(
-        self, 
-        inputs: tuple[tf.Tensor, tf.Tensor], 
-        mask: bool | None = None
-    ) -> tf.Tensor | None:
-        # if self.mask_value is None:
-        #     return None
-        _, indices = inputs
-        return keras.ops.not_equal(indices, self.mask_value)
-    
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class AminoAcidType(descriptors.Descriptor):
-
-    def __init__(self, vocab=None, **kwargs):
-        vocab = [
-            "A", "C", "D", "E", "F", "G", "H", "I", "K", "L", 
-            "M", "N", "P", "Q", "R", "S", "T", "V", "W", "Y",
-        ]
-        super().__init__(vocab=vocab, **kwargs)
-
-    def call(self, mol: chem.Mol) -> list[str]:
-        residue = residues_reverse.get(mol.canonical_smiles)
-        if not residue:
-            raise KeyError(f'Could not find {mol.canonical_smiles} in `residues_reverse`.')
-        mol = chem.remove_hs(mol)
-        return _extract_residue_type(residues_reverse[mol.canonical_smiles])
-    
-def sequence_split(sequence: str):
-    patterns = [
-        r'(\[[A-Za-z0-9]+\]-[A-Z]\[[A-Za-z0-9]+\])', # N-term mod + mod
-        r'([A-Z]\[[A-Za-z0-9]+\]-\[[A-Za-z0-9]+\])', # C-term mod + mod
-        r'([A-Z]-\[[A-Za-z0-9]+\])', # C-term mod
-        r'(\[[A-Za-z0-9]+\]-[A-Z])', # N-term mod
-        r'([A-Z]\[[A-Za-z0-9]+\])', # Mod
-        r'([A-Z])', # No mod
-    ]
-    return [match.group(0) for match in re.finditer("|".join(patterns), sequence)]
-
-residues = {
-    "A": "N[C@@H](C)C(=O)O",
-    "C": "N[C@@H](CS)C(=O)O",
-    "C[Carbamidomethyl]": "N[C@@H](CSCC(=O)N)C(=O)O",
-    "D": "N[C@@H](CC(=O)O)C(=O)O",
-    "E": "N[C@@H](CCC(=O)O)C(=O)O",
-    "F": "N[C@@H](Cc1ccccc1)C(=O)O",
-    "G": "NCC(=O)O",
-    "H": "N[C@@H](CC1=CN=C-N1)C(=O)O",
-    "I": "N[C@@H](C(CC)C)C(=O)O",
-    "K": "N[C@@H](CCCCN)C(=O)O",
-    "K[Acetyl]": "N[C@@H](CCCCNC(=O)C)C(=O)O",
-    "K[Crotonyl]": "N[C@@H](CCCCNC(C=CC)=O)C(=O)O",
-    "K[Dimethyl]": "N[C@@H](CCCCN(C)C)C(=O)O",
-    "K[Formyl]": "N[C@@H](CCCCNC=O)C(=O)O",
-    "K[Malonyl]": "N[C@@H](CCCCNC(=O)CC(O)=O)C(=O)O",
-    "K[Methyl]": "N[C@@H](CCCCNC)C(=O)O",
-    "K[Propionyl]": "N[C@@H](CCCCNC(=O)CC)C(=O)O",
-    "K[Succinyl]": "N[C@@H](CCCCNC(CCC(O)=O)=O)C(=O)O",
-    "K[Trimethyl]": "N[C@@H](CCCC[N+](C)(C)C)C(=O)O",
-    "L": "N[C@@H](CC(C)C)C(=O)O",
-    "M": "N[C@@H](CCSC)C(=O)O",
-    "M[Oxidation]": "N[C@@H](CCS(=O)C)C(=O)O",
-    "N": "N[C@@H](CC(=O)N)C(=O)O",
-    "P": "N1[C@@H](CCC1)C(=O)O",
-    "P[Oxidation]": "N1CC(O)C[C@H]1C(=O)O",
-    "Q": "N[C@@H](CCC(=O)N)C(=O)O",
-    "R": "N[C@@H](CCCNC(=N)N)C(=O)O",
-    "R[Deamidated]": "N[C@@H](CCCNC(N)=O)C(=O)O",
-    "R[Dimethyl]": "N[C@@H](CCCNC(N(C)C)=N)C(=O)O",
-    "R[Methyl]": "N[C@@H](CCCNC(=N)NC)C(=O)O",
-    "S": "N[C@@H](CO)C(=O)O",
-    "T": "N[C@@H](C(O)C)C(=O)O",
-    "V": "N[C@@H](C(C)C)C(=O)O",
-    "W": "N[C@@H](CC(=CN2)C1=C2C=CC=C1)C(=O)O",
-    "Y": "N[C@@H](Cc1ccc(O)cc1)C(=O)O",
-    "Y[Nitro]": "N[C@@H](Cc1ccc(O)c(N(=O)=O)c1)C(=O)O",
-    "Y[Phospho]": "N[C@@H](Cc1ccc(OP(O)(=O)O)cc1)C(=O)O",
-    "[Acetyl]-A": "N(C(C)=O)[C@@H](C)C(=O)O",
-    "[Acetyl]-C": "N(C(C)=O)[C@@H](CS)C(=O)O",
-    "[Acetyl]-D": "N(C(=O)C)[C@H](C(=O)O)CC(=O)O",
-    "[Acetyl]-E": "N(C(=O)C)[C@@H](CCC(O)=O)C(=O)O",
-    "[Acetyl]-F": "N(C(C)=O)[C@@H](Cc1ccccc1)C(=O)O",
-    "[Acetyl]-G": "N(C(=O)C)CC(=O)O",
-    "[Acetyl]-H": "N(C(=O)C)[C@@H](Cc1[nH]cnc1)C(=O)O",
-    "[Acetyl]-I": "N(C(=O)C)[C@@H]([C@H](CC)C)C(=O)O",
-    "[Acetyl]-K": "N(C(C)=O)[C@@H](CCCCN)C(=O)O",
-    "[Acetyl]-L": "N(C(=O)C)[C@@H](CC(C)C)C(=O)O",
-    "[Acetyl]-M": "N(C(=O)C)[C@@H](CCSC)C(=O)O",
-    "[Acetyl]-N": "N(C(C)=O)[C@@H](CC(=O)N)C(=O)O",
-    "[Acetyl]-P": "N1(C(=O)C)CCC[C@H]1C(=O)O",
-    "[Acetyl]-Q": "N(C(=O)C)[C@@H](CCC(=O)N)C(=O)O",
-    "[Acetyl]-R": "N(C(C)=O)[C@@H](CCCN=C(N)N)C(=O)O",
-    "[Acetyl]-S": "N(C(C)=O)[C@@H](CO)C(=O)O",
-    "[Acetyl]-T": "N(C(=O)C)[C@@H]([C@H](O)C)C(=O)O",
-    "[Acetyl]-V": "N(C(=O)C)[C@@H](C(C)C)C(=O)O",
-    "[Acetyl]-W": "N(C(C)=O)[C@@H](Cc1c2ccccc2[nH]c1)C(=O)O",
-    "[Acetyl]-Y": "N(C(C)=O)[C@@H](Cc1ccc(O)cc1)C(=O)O"
-}
-
-residues_reverse = {}
-def register_peptide_residues(residues: dict[str, str]):
-    for residue, smiles in residues.items():
-        residues[residue] = Chem.MolToSmiles(Chem.MolFromSmiles(smiles))
-        residues_reverse[residues[residue]] = residue
-    
-register_peptide_residues(residues)
-
-def _extract_residue_type(residue_tag: str) -> str:
-    pattern = r"(?<!\[)[A-Z](?![\w-])"
-    return [match.group(0) for match in re.finditer(pattern, residue_tag)][0]

molcraft-0.1.0a3.dist-info/RECORD

@@ -1,20 +0,0 @@
-molcraft/__init__.py,sha256=2ZNfWBjGl8DscOwjdDiRkgIsuPnKit29Q3MhZyP336Q,435
-molcraft/callbacks.py,sha256=6gwCwdsHGb-fVB4m1QGmtBwQwZ9mFq9QUkmPKSMn05U,849
-molcraft/chem.py,sha256=_UO5O-I7KUtGf3vRrFEYoAUGlW5xi2x8ylu5f-Ybumo,18696
-molcraft/conformers.py,sha256=p09gOQOdxLSj3yohZOMkxxLriHsZ1ZqOoiWLi73OpIg,4325
-molcraft/datasets.py,sha256=rFgXTC1ZheLhfgQgcCspP_wEE54a33PIneH7OplbS-8,4047
-molcraft/descriptors.py,sha256=x6RfZ-gK7D_WSvmK6sh6yHyEjQqovPnRU0xwC3dAKfg,2880
-molcraft/features.py,sha256=69oV_GHNdBKPA4sp6Tpo6brvNmaauk_IVIzNjX7VDmg,13648
-molcraft/featurizers.py,sha256=Yu8I6I_zkzB__WYSiqz-FDGjvKFOmyWFxojRBr39Aw8,26236
-molcraft/layers.py,sha256=HjnAtqhuP0uZ5yP4L33k3xT4IUdLavWBrjd3wO9_Rmw,64915
-molcraft/models.py,sha256=DXqWR_XnMVXQseVR91XnDLXvmHa1hv-6_Y_wvpQZBFI,17476
-molcraft/ops.py,sha256=iiE6zgA2P7cmjKO1RHmL9GE_Tv7Tyuo_xDoxB_ELZQM,3824
-molcraft/records.py,sha256=w4-bcWZEC0oVInrE1e0kQBroIaSCA0PN1JBPOtO6VUY,5251
-molcraft/tensors.py,sha256=b7PO-YOvV72s9g057ILJACKS2n2fn10VkO35gHXpssI,22312
-molcraft/experimental/__init__.py,sha256=x5h6LOO8bo3NPjkKKM9M1H-Kz6R3yxYhRSePoxHCdRE,42
-molcraft/experimental/peptides.py,sha256=RCuOTSwoYHGSdeYi6TWHdPIv2WC3avCZjKLdhEZQeXw,8997
-molcraft-0.1.0a3.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
-molcraft-0.1.0a3.dist-info/METADATA,sha256=f_5sBinpFcGSqKLaSqGkZJ83gGQtZw1Pb3fkgq9aCBM,4088
-molcraft-0.1.0a3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-molcraft-0.1.0a3.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
-molcraft-0.1.0a3.dist-info/RECORD,,