molcraft 0.1.0a21__py3-none-any.whl → 0.1.0a22__py3-none-any.whl

molcraft/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.0a21'
+__version__ = '0.1.0a22'
 
 import os
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -14,6 +14,4 @@ from molcraft import records
 from molcraft import tensors
 from molcraft import callbacks
 from molcraft import datasets
-from molcraft import losses
-
-from molcraft.applications import proteomics
+from molcraft import losses

molcraft/chem.py

@@ -22,12 +22,19 @@ class Mol(Chem.Mol):
         if explicit_hs:
             rdkit_mol = Chem.AddHs(rdkit_mol) 
         rdkit_mol.__class__ = cls 
+        setattr(rdkit_mol, '_encoding', encoding)
         return rdkit_mol
 
     @property
     def canonical_smiles(self) -> str:
         return Chem.MolToSmiles(self, canonical=True)
     
+    @property 
+    def encoding(self):
+        if hasattr(self, '_encoding'):
+            return self._encoding 
+        return None 
+    
     @property
     def bonds(self) -> list['Bond']:
         if not hasattr(self, '_bonds'):
@@ -391,6 +398,7 @@ def embed_conformers(
     mol: Mol, 
     num_conformers: int, 
     method: str = 'ETKDGv3',
+    timeout: int | None = None, 
     random_seed: int | None = None, 
     **kwargs
 ) -> Mol:
@@ -403,6 +411,7 @@ def embed_conformers(
         'KDG': rdDistGeom.KDG()
     }
     mol = Mol(mol)
+    encoding = mol.encoding or mol.canonical_smiles
     embedding_method = available_embedding_methods.get(method)
     if embedding_method is None:
         raise ValueError(
@@ -413,8 +422,14 @@ def embed_conformers(
     for key, value in kwargs.items():
         setattr(embedding_method, key, value)
 
-    if random_seed is not None:
-        embedding_method.randomSeed = random_seed
+    if not timeout:
+        timeout = 0 # No timeout
+
+    if not random_seed:
+        random_seed = -1 # No random seed
+
+    embedding_method.randomSeed = random_seed
+    embedding_method.timeout = timeout
 
     success = rdDistGeom.EmbedMultipleConfs(
         mol, numConfs=num_conformers, params=embedding_method
@@ -422,17 +437,18 @@ def embed_conformers(
     num_successes = len(success)
     if num_successes < num_conformers:
         warnings.warn(
-            f'Could only embed {num_successes} out of {num_conformers} conformer(s) '
-            f'for {mol.canonical_smiles!r} using {method}. Embedding the remaining '
-            f'{num_conformers - num_successes} conformer(s) using different embedding methods.',
-            stacklevel=2
+            f'Could only embed {num_successes} out of {num_conformers} conformer(s) for '
+            f'{encoding!r} using the specified method ({method!r}) and parameters. Attempting '
+            f'to embed the remaining {num_conformers-num_successes} using fallback methods.',
         )
+        max_iters = 20 * mol.num_atoms # Doubling the number of iterations
         for fallback_method in [method, 'ETDG', 'KDG']:
             fallback_embedding_method = available_embedding_methods[fallback_method]
             fallback_embedding_method.useRandomCoords = True
+            fallback_embedding_method.maxIterations = int(max_iters)
             fallback_embedding_method.clearConfs = False
-            if random_seed is not None:
-                fallback_embedding_method.randomSeed = random_seed
+            fallback_embedding_method.timeout = int(timeout)
+            fallback_embedding_method.randomSeed = int(random_seed)
             success = rdDistGeom.EmbedMultipleConfs(
                 mol, numConfs=(num_conformers - num_successes), params=fallback_embedding_method
             )
@@ -441,7 +457,7 @@ def embed_conformers(
                 break
         else:
             raise RuntimeError(
-                f'Could not embed {num_conformers} conformer(s) for {mol.canonical_smiles!r}. '
+                f'Could not embed {num_conformers} conformer(s) for {encoding!r}. '
             )
     return mol
 
@@ -485,9 +501,8 @@ def optimize_conformers(
     except RuntimeError as e:
         warnings.warn(
             f'{method} force field minimization did not succeed. Proceeding without it.',
-            stacklevel=2
         )
-        return mol
+        return Mol(mol)
     return mol_optimized
 
 def prune_conformers(
@@ -500,7 +515,6 @@ def prune_conformers(
         warnings.warn(
             'Molecule has no conformers. To embed conformers, invoke the `embed` method, '
             'and optionally followed by `minimize()` to perform force field minimization.',
-            stacklevel=2
         )
         return mol
     

molcraft/featurizers.py

@@ -2,8 +2,7 @@ import keras
 import json
 import abc
 import typing 
-import copy 
-import warnings
+import os
 import numpy as np
 import pandas as pd
 import tensorflow as tf
@@ -13,6 +12,7 @@ from pathlib import Path
 
 from molcraft import tensors 
 from molcraft import features
+from molcraft import records
 from molcraft import chem
 from molcraft import descriptors
 
@@ -41,6 +41,17 @@ class GraphFeaturizer(abc.ABC):
     def load(filepath: str | Path, *args, **kwargs) -> 'GraphFeaturizer':
         return load_featurizer(filepath, *args, **kwargs)
     
+    def write_records(self, inputs: str | chem.Mol | tuple, path: str | Path, **kwargs) -> None:
+         records.write(
+            inputs, featurizer=self, path=path, **kwargs
+         )
+
+    @staticmethod
+    def read_records(path: str | Path, **kwargs) -> tf.data.Dataset:
+        return records.read(
+            path=path, **kwargs
+        )
+        
     def __call__(
         self,
         inputs: str | chem.Mol | tuple | typing.Iterable,

molcraft/records.py

@@ -3,20 +3,24 @@ import math
 import glob
 import time
 import typing 
+import warnings
 import tensorflow as tf
 import numpy as np
 import pandas as pd
 import multiprocessing as mp
 
 from molcraft import tensors
-from molcraft import featurizers
+
+if typing.TYPE_CHECKING:
+    from molcraft import featurizers
 
 
 def write(
     inputs: list[str | tuple], 
-    featurizer: featurizers.GraphFeaturizer,
+    featurizer: 'featurizers.GraphFeaturizer',
     path: str, 
-    overwrite: bool = True, 
+    exist_ok: bool = False,
+    overwrite: bool = False, 
     num_files: typing.Optional[int] = None, 
     num_processes: typing.Optional[int] = None,
     multiprocessing: bool = False,
@@ -24,6 +28,8 @@ def write(
 ) -> None:
     
     if os.path.isdir(path):
+        if not exist_ok:
+            raise FileExistsError(f'Records already exist: {path}')
         if not overwrite:
             return
         else:
@@ -60,9 +66,11 @@ def write(
             chunk_sizes[i % num_files] += 1
         
         input_chunks = []
+        start_indices = []
         current_index = 0
         for size in chunk_sizes:
             input_chunks.append(inputs[current_index: current_index + size])
+            start_indices.append(current_index)
             current_index += size 
         
         assert current_index == num_examples
@@ -73,13 +81,13 @@ def write(
         ]
         
         if not multiprocessing:
-            for path, input_chunk in zip(paths, input_chunks):
-                _write_tfrecord(input_chunk, path, featurizer)
+            for path, input_chunk, start_index in zip(paths, input_chunks, start_indices):
+                _write_tfrecord(input_chunk, path, featurizer, start_index)
             return
         
         processes = []
         
-        for path, input_chunk in zip(paths, input_chunks):
+        for path, input_chunk, start_index in zip(paths, input_chunks, start_indices):
         
             while len(processes) >= num_processes:
                 for process in processes:
@@ -91,7 +99,7 @@ def write(
                     
             process = mp.Process(
                 target=_write_tfrecord,
-                args=(input_chunk, path, featurizer)
+                args=(input_chunk, path, featurizer, start_index)
             )
             processes.append(process)
             process.start()
@@ -134,9 +142,10 @@ def load_spec(path: str) -> tensors.GraphTensor.Spec:
     return spec
     
 def _write_tfrecord(
-    inputs, 
+    inputs: list[str, tuple], 
     path: str,
-    featurizer: featurizers.GraphFeaturizer, 
+    featurizer: 'featurizers.GraphFeaturizer', 
+    start_index: int,
 ) -> None:
     
     def _write_example(tensor):
@@ -147,12 +156,17 @@ def _write_tfrecord(
         writer.write(serialized_feature)
 
     with tf.io.TFRecordWriter(path) as writer:
-        for x in inputs:
+        for i, x in enumerate(inputs):
             if isinstance(x, (list, np.ndarray)):
                 x = tuple(x)
-            tensor = featurizer(x)
-            if tensor is not None:
+            try:
+                tensor = featurizer(x)
                 _write_example(tensor)
+            except Exception as e:
+                warnings.warn(
+                    f"Could not write record for index {i + start_index}, proceeding without it."
+                    f"Exception raised:\n{e}"
+                )
 
 def _serialize_example(
     feature: dict[str, tf.train.Feature]

{molcraft-0.1.0a21.dist-info → molcraft-0.1.0a22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a21
+Version: 0.1.0a22
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License

{molcraft-0.1.0a21.dist-info → molcraft-0.1.0a22.dist-info}/RECORD

@@ -1,21 +1,21 @@
-molcraft/__init__.py,sha256=z-THS82UDJrfPipyQx73udsMNMNVb0ah5X4sfxGLk6M,477
+molcraft/__init__.py,sha256=O88EmicQAD8oz9oFMXk_IzFChQEbbU-BCs3IE-c9Dkk,431
 molcraft/callbacks.py,sha256=x5HnkZhqcFRrW6xdApt_jZ4X08A-0fxcnFKfdmRKa0c,3571
-molcraft/chem.py,sha256=NvjwAAEy6RtbUICb3fKKCR1pQELrVWwrrkkUL_mazsU,21534
+molcraft/chem.py,sha256=ynrEpWZL2D370p7CqH2kE1KhBByq7IiuQbUNoKQt96I,22028
 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=b9b8XUGqMku3EwwQcjeE2rsaMQ0T_XovupUXjX0Awug,17724
+molcraft/featurizers.py,sha256=1yBz5-JA7IhNm0dGivvVm1nJ5QGck8VQXtwHPWFbTuQ,18091
 molcraft/layers.py,sha256=H7XZru4XGJA6gbRO9V1BsGqh1mIrMdhzNCKS5o6oNok,64544
 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/records.py,sha256=sopYElKWC3A9QE5I8_957v3faLb2Wt5WILHZv_FLLds,6283
 molcraft/tensors.py,sha256=vk-W8zZu-re1g18YevDEEoVQRxT4AdIiMdI-4EvtJI4,22477
 molcraft/applications/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 molcraft/applications/chromatography.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 molcraft/applications/proteomics.py,sha256=BL3EtW-q-0j79pLYO7npC67mA2ApRhH-XI4rOaP8_wc,8407
-molcraft-0.1.0a21.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
-molcraft-0.1.0a21.dist-info/METADATA,sha256=PL2qoRwua1MojNjPA4WDHJY1BkIjcvjOeCxNvoT44KA,3930
-molcraft-0.1.0a21.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-molcraft-0.1.0a21.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
-molcraft-0.1.0a21.dist-info/RECORD,,
+molcraft-0.1.0a22.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a22.dist-info/METADATA,sha256=1OHx3-Q94fFEi21l0p3bnMjU-Q0EHaZLm4PU1A6QbkU,3930
+molcraft-0.1.0a22.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+molcraft-0.1.0a22.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a22.dist-info/RECORD,,