molcraft 0.1.0a16__py3-none-any.whl → 0.1.0a17__py3-none-any.whl

molcraft/layers.py

@@ -1441,7 +1441,7 @@ class AddContext(GraphLayer):
     def __init__(
         self, 
         field: str = 'feature',
-        drop: bool = True,
+        drop: bool = False,
         normalize: bool = False,
         **kwargs
     ) -> None:

molcraft/models.py

@@ -154,6 +154,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
         return graph
     
     def get_config(self):
+        """Obtain model config."""
         config = super().get_config()
         if hasattr(self, '_model_layers') and self._model_layers is not None:
             config['model_layers'] = [
@@ -164,6 +165,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
     
     @classmethod
     def from_config(cls, config: dict):
+        """Obtain model from model config."""
         if 'model_layers' in config:
             config['model_layers'] = [
                 keras.saving.deserialize_keras_object(l) 

molcraft/records.py

@@ -14,7 +14,7 @@ from molcraft import featurizers
 
 def write(
     inputs: list[str | tuple], 
-    featurizer: featurizers.Featurizer,
+    featurizer: featurizers.GraphFeaturizer,
     path: str, 
     overwrite: bool = True, 
     num_files: typing.Optional[int] = None, 
@@ -23,10 +23,13 @@ def write(
     device: str = '/cpu:0'
 ) -> None:
     
-    if os.path.isdir(path) and not overwrite:
-        return 
-    
-    os.makedirs(path, exist_ok=True)
+    if os.path.isdir(path):
+        if not overwrite:
+            return
+        else:
+            _remove_files(path)
+    else:
+        os.makedirs(path)
 
     with tf.device(device):
 
@@ -133,7 +136,7 @@ def load_spec(path: str) -> tensors.GraphTensor.Spec:
 def _write_tfrecord(
     inputs, 
     path: str,
-    featurizer: featurizers.Featurizer, 
+    featurizer: featurizers.GraphFeaturizer, 
 ) -> None:
     
     def _write_example(tensor):
@@ -149,11 +152,7 @@ def _write_tfrecord(
                 x = tuple(x)
             tensor = featurizer(x)
             if tensor is not None:
-                if isinstance(tensor, tensors.GraphTensor):
-                    _write_example(tensor)
-                else:
-                    for t in tensor:
-                        _write_example(t)
+                _write_example(tensor)
 
 def _serialize_example(
     feature: dict[str, tf.train.Feature]
@@ -168,8 +167,18 @@ def _parse_example(
 ) -> tf.Tensor:
     out = tf.io.parse_single_example(
         x, features={'feature': tf.io.RaggedFeature(tf.string)})['feature']
-    out = [tf.ensure_shape(tf.io.parse_tensor(x[0], s.dtype), s.shape) for (x, s) in zip(
-        tf.split(out, len(tf.nest.flatten(spec, expand_composites=True))), 
-        tf.nest.flatten(spec, expand_composites=True))]
+    out = [
+        tf.ensure_shape(tf.io.parse_tensor(x[0], s.dtype), s.shape) 
+        for (x, s) in zip(
+            tf.split(out, len(tf.nest.flatten(spec, expand_composites=True))), 
+            tf.nest.flatten(spec, expand_composites=True)
+        )
+    ]
     out = tf.nest.pack_sequence_as(spec, tf.nest.flatten(out), expand_composites=True)
-    return out
+    return out
+
+def _remove_files(path):
+    for filename in os.listdir(path):
+        if filename.endswith('tfrecord') or filename == 'spec.pb':
+            filepath = os.path.join(path, filename)
+            os.remove(filepath)

{molcraft-0.1.0a16.dist-info → molcraft-0.1.0a17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a16
+Version: 0.1.0a17
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -43,9 +43,9 @@ Provides-Extra: gpu
 Requires-Dist: tensorflow[and-cuda]>=2.16; extra == "gpu"
 Dynamic: license-file
 
-<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo">
+<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo", width="90%">
 
-**Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML. 
+**Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML.
 
 > [!NOTE]  
 > In progress.
@@ -83,11 +83,12 @@ featurizer = featurizers.MolGraphFeaturizer(
         features.BondType(),
         features.IsRotatable(),
     ],
-    super_atom=True,
+    super_node=True,
     self_loops=True,
+    include_hydrogens=False,
 )
 
-graph = featurizer([('N[C@@H](C)C(=O)O', 2.0), ('N[C@@H](CS)C(=O)O', 1.0)])
+graph = featurizer([('N[C@@H](C)C(=O)O', 2.5), ('N[C@@H](CS)C(=O)O', 1.5)])
 print(graph)
 
 model = models.GraphModel.from_layers(
@@ -95,13 +96,13 @@ model = models.GraphModel.from_layers(
         layers.Input(graph.spec),
         layers.NodeEmbedding(dim=128),
         layers.EdgeEmbedding(dim=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.Readout(mode='mean'),
-        keras.layers.Dense(units=1024, activation='relu'),
-        keras.layers.Dense(units=1024, activation='relu'),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.Readout(),
+        keras.layers.Dense(units=1024, activation='elu'),
+        keras.layers.Dense(units=1024, activation='elu'),
         keras.layers.Dense(1)
     ]
 )

molcraft-0.1.0a17.dist-info/RECORD

@@ -0,0 +1,21 @@
+molcraft/__init__.py,sha256=vc-z1sgtzPY7Spwzkemu7I_b9ekEN9egnHrLEKbB9bk,431
+molcraft/callbacks.py,sha256=x5HnkZhqcFRrW6xdApt_jZ4X08A-0fxcnFKfdmRKa0c,3571
+molcraft/chem.py,sha256=e56qBDuqh8rq_4-UMyp6LCQNxxSx8hZ7gzuz-87DHgw,21652
+molcraft/datasets.py,sha256=Nd2lw5USUZE52vvAiNr-q-n03Y3--NlZlK0NzqHgp-E,4145
+molcraft/descriptors.py,sha256=Cl3KnBPsTST7XLgRLktkX5LwY9MV0P_lUlrt8iPV5no,3508
+molcraft/features.py,sha256=s0WeV8eZcDEypPgC1m37f4s9QkvWIlVgn-L43Cdsa14,13525
+molcraft/featurizers.py,sha256=bD3RFY9eg89-O-Nxgy6gote1zS4cyjOgzdSiSJZJdJE,17664
+molcraft/layers.py,sha256=Y-TMb4oHh3R7tHgr7f3Y8sEPDnoSTbtwB6NkZIVnmcA,61734
+molcraft/losses.py,sha256=qnS2yC5g-O3n_zVea9MR6TNiFraW2yqRgePOisoUP4A,1065
+molcraft/models.py,sha256=2Pc1htT9fCukGd8ZxrvE0rzEHsPBm0pluHw4FZXaUE4,21963
+molcraft/ops.py,sha256=bQbdFDt9waxVCzF5-dkTB6vlpj9eoSt8I4Qg7ZGXbsU,6178
+molcraft/records.py,sha256=0j4EWP55sfnkoQIH5trdaAIevPfVbAtPLrygTRmLyFw,5686
+molcraft/tensors.py,sha256=EOUKx496KUZsjA1zA2ABc7tU_TW3Jv7AXDsug_QsLbA,22407
+molcraft/applications/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+molcraft/applications/chromatography.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+molcraft/applications/proteomics.py,sha256=Jb7OwHJHc_I7Wk3qnqr40j9P7um2EKtUnB4r-XhrnAc,7180
+molcraft-0.1.0a17.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a17.dist-info/METADATA,sha256=XqNJDwFfY6pWNqQKYLyUOxwyvmfYUkOWTKou-ZQYXL4,3930
+molcraft-0.1.0a17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+molcraft-0.1.0a17.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a17.dist-info/RECORD,,

molcraft/conformers.py

@@ -1,151 +0,0 @@
-import keras
-
-from molcraft import chem
-
-
-@keras.saving.register_keras_serializable(package="molcraft")
-class ConformerProcessor:
-
-    def get_config(self) -> dict:
-        return {}
-
-    @classmethod
-    def from_config(cls, config: dict):
-        return cls(**config)
-
-    def __call__(self, mol: chem.Mol) -> chem.Mol:
-        raise NotImplementedError
-
-
-@keras.saving.register_keras_serializable(package="molcraft")
-class ConformerEmbedder(ConformerProcessor):
-
-    def __init__(
-        self, 
-        method: str = 'ETKDGv3',
-        num_conformers: int = 5, 
-        **kwargs,
-    ) -> None:
-        self.method = method 
-        self.num_conformers = num_conformers 
-        self.kwargs = kwargs 
-
-    def get_config(self) -> dict:
-        config = {
-            'method': self.method, 
-            'num_conformers': self.num_conformers, 
-        }
-        config.update({
-            k: v for (k, v) in self.kwargs.items()
-        })
-        return config
-    
-    def __call__(self, mol: chem.Mol) -> chem.Mol:
-        return chem.embed_conformers(
-            mol, 
-            method=self.method,
-            num_conformers=self.num_conformers,
-            **self.kwargs,
-        )
-
-
-@keras.saving.register_keras_serializable(package="molcraft")
-class ConformerOptimizer(ConformerProcessor):
-
-    def __init__(
-        self, 
-        method: str = 'UFF',
-        max_iter: int = 200, 
-        ignore_interfragment_interactions: bool = True,
-        vdw_threshold: float = 10.0,
-        **kwargs,
-    ) -> None:
-        self.method = method 
-        self.max_iter = max_iter 
-        self.ignore_interfragment_interactions = ignore_interfragment_interactions
-        self.vdw_threshold = vdw_threshold 
-        self.kwargs = kwargs
-
-    def get_config(self) -> dict:
-        config = {
-            'method': self.method,
-            'max_iter': self.max_iter,
-            'ignore_interfragment_interactions': self.ignore_interfragment_interactions,
-            'vdw_threshold': self.vdw_threshold,
-        }
-        config.update({
-            k: v for (k, v) in self.kwargs.items()
-        })
-        return config
-    
-    def __call__(self, mol: chem.Mol) -> chem.Mol:
-        return chem.optimize_conformers(
-            mol,
-            method=self.method,
-            max_iter=self.max_iter,
-            ignore_interfragment_interactions=self.ignore_interfragment_interactions,
-            vdw_threshold=self.vdw_threshold,
-            **self.kwargs,
-        )
-
-
-@keras.saving.register_keras_serializable(package="molcraft")
-class ConformerPruner(ConformerProcessor):
-    def __init__(
-        self,
-        keep: int = 1,
-        threshold: float = 0.0,
-        energy_force_field: str = 'UFF',
-        **kwargs,
-    ) -> None:
-        self.keep = keep
-        self.threshold = threshold
-        self.energy_force_field = energy_force_field
-        self.kwargs = kwargs
-
-    def get_config(self) -> dict:
-        config = {
-            'keep': self.keep,
-            'threshold': self.threshold,
-            'energy_force_field': self.energy_force_field,
-        }
-        config.update({
-            k: v for (k, v) in self.kwargs.items()
-        })
-        return config
-    
-    def __call__(self, mol: chem.Mol) -> chem.Mol:
-        return chem.prune_conformers(
-            mol,
-            keep=self.keep,
-            threshold=self.threshold,
-            energy_force_field=self.energy_force_field,
-            **self.kwargs,
-        )
-
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class ConformerGenerator(ConformerProcessor):
-
-    def __init__(self, steps: list[ConformerProcessor]) -> None:
-        self.steps = steps
-
-    def get_config(self) -> dict:
-        return {
-            "steps": [
-                keras.saving.serialize_keras_object(step) for step in self.steps
-            ]
-        }
-
-    @classmethod
-    def from_config(cls, config: dict) -> 'ConformerGenerator':
-        steps = [
-            keras.saving.deserialize_keras_object(obj) 
-            for obj in config["steps"]
-        ]
-        return cls(steps)
-
-    def __call__(self, mol: chem.Mol) -> chem.Mol:
-        for step in self.steps:
-            mol = step(mol)
-        return mol

molcraft-0.1.0a16.dist-info/RECORD

@@ -1,21 +0,0 @@
-molcraft/__init__.py,sha256=uo2ze7WMv3VhP0JcJedXDS9UkeEFydlXgSw4l7Xi8E0,464
-molcraft/callbacks.py,sha256=x5HnkZhqcFRrW6xdApt_jZ4X08A-0fxcnFKfdmRKa0c,3571
-molcraft/chem.py,sha256=--4AdZV0TCj_cf5i-TRidNJGSFyab1ksUEMjmDi7zaM,21837
-molcraft/conformers.py,sha256=K6ZtiSUNDN_fwqGP9JrPcwALLFFvlMlF_XejEJH3Sr4,4205
-molcraft/datasets.py,sha256=QKHi9SUBKvJvdkRFmRQNowhrnu35pQqtujuLatOK8bE,4151
-molcraft/descriptors.py,sha256=jJpT0XWu3Tx_bxnwk1rENySRkaM8cMDMaDIjG8KKvtg,3097
-molcraft/features.py,sha256=GwOecLCNUIuGfbIVzsAJH4LikkzWMKj5IT7zSgGTttU,13846
-molcraft/featurizers.py,sha256=8Jmd2yguYmVRyh5wkn6sRzzEENkJ0TqHSlR8qgC4zNY,27131
-molcraft/layers.py,sha256=200Y4QLOXDyHw1bnjoSQ6hZ-zD0vpforv-KQGESAZi8,61733
-molcraft/losses.py,sha256=qnS2yC5g-O3n_zVea9MR6TNiFraW2yqRgePOisoUP4A,1065
-molcraft/models.py,sha256=hKYSV8z65ohRKfPyjjzxZeVjipm064BWeUBGZE0tpyU,21882
-molcraft/ops.py,sha256=bQbdFDt9waxVCzF5-dkTB6vlpj9eoSt8I4Qg7ZGXbsU,6178
-molcraft/records.py,sha256=MbvYkcCunbAmpy_MWXmQ9WBGi2WvwxFUlwQSPKPvSSk,5534
-molcraft/tensors.py,sha256=EOUKx496KUZsjA1zA2ABc7tU_TW3Jv7AXDsug_QsLbA,22407
-molcraft/applications/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-molcraft/applications/proteomics.py,sha256=usZkoYtmTi1BtoP8SigyBNPjxR-nLH1yEsuAdpjvF2M,9009
-molcraft-0.1.0a16.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
-molcraft-0.1.0a16.dist-info/METADATA,sha256=JL32RxGZY92s39EnhonY4mZbvgGvXjcybtV9nXukn4s,3930
-molcraft-0.1.0a16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-molcraft-0.1.0a16.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
-molcraft-0.1.0a16.dist-info/RECORD,,