molcraft 0.1.0rc10__py3-none-any.whl

molcraft/ops.py

@@ -0,0 +1,195 @@
+import warnings
+import keras
+import numpy as np
+import tensorflow as tf
+from keras import backend
+
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def gather(
+    node_feature: tf.Tensor, 
+    edge: tf.Tensor
+) -> tf.Tensor:
+    if backend.backend() == 'tensorflow':
+        return tf.gather(node_feature, edge)
+    expected_rank = len(keras.ops.shape(node_feature))
+    current_rank = len(keras.ops.shape(edge))
+    for _ in range(expected_rank - current_rank):
+        edge = keras.ops.expand_dims(edge, axis=-1)
+    return keras.ops.take_along_axis(node_feature, edge, axis=0)
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def aggregate(
+    node_feature: tf.Tensor, 
+    edge: tf.Tensor, 
+    num_nodes: tf.Tensor,
+    mode: str = 'sum',
+) -> tf.Tensor:
+    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
+    )
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def propagate(
+    node_feature: tf.Tensor,
+    edge_source: tf.Tensor,
+    edge_target: tf.Tensor,
+    edge_feature: tf.Tensor | None = None,
+    edge_weight: tf.Tensor | None = None,
+) -> tf.Tensor:
+    num_nodes = keras.ops.shape(node_feature)[0]
+
+    node_feature_source = gather(node_feature, edge_source)
+    
+    if edge_weight is not None:
+        node_feature_source *= edge_weight
+
+    if edge_feature is not None:
+        node_feature_source += edge_feature
+        
+    return aggregate(node_feature, edge_target, num_nodes)
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def scatter_update(
+    inputs: tf.Tensor,
+    indices: tf.Tensor,
+    updates: tf.Tensor,
+) -> tf.Tensor:
+    if indices.dtype == tf.bool:
+        indices = keras.ops.stack(keras.ops.where(indices), axis=-1)
+    expected_rank = len(keras.ops.shape(inputs))
+    current_rank = len(keras.ops.shape(indices))
+    for _ in range(expected_rank - current_rank):
+        indices = keras.ops.expand_dims(indices, axis=-1)
+    return keras.ops.scatter_update(inputs, indices, updates)
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def scatter_add(
+    inputs: tf.Tensor,
+    indices: tf.Tensor,
+    updates: tf.Tensor,
+) -> tf.Tensor:
+    if indices.dtype == tf.bool:
+        indices = keras.ops.stack(keras.ops.where(indices), axis=-1)
+    expected_rank = len(keras.ops.shape(inputs))
+    current_rank = len(keras.ops.shape(indices))
+    for _ in range(expected_rank - current_rank):
+        indices = keras.ops.expand_dims(indices, axis=-1)
+    if backend.backend() == 'tensorflow':
+        return tf.tensor_scatter_nd_add(inputs, indices, updates)
+    updates = scatter_update(keras.ops.zeros_like(inputs), indices, updates)
+    return inputs + updates
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def edge_softmax(
+    score: tf.Tensor, 
+    edge_target: tf.Tensor
+) -> tf.Tensor:
+    num_segments = keras.ops.cond(
+        keras.ops.greater(keras.ops.shape(edge_target)[0], 0),
+        lambda: keras.ops.maximum(keras.ops.max(edge_target) + 1, 1),
+        lambda: 0
+    )
+    score_max = keras.ops.segment_max(
+        score, edge_target, num_segments, sorted=False
+    )
+    score_max = gather(score_max, edge_target)
+    numerator = keras.ops.exp(score - score_max)
+    denominator = keras.ops.segment_sum(
+        numerator, edge_target, num_segments, sorted=False
+    )
+    denominator = gather(denominator, edge_target)
+    return numerator / denominator
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def edge_weight(
+    edge: tf.Tensor,
+    edge_weight: tf.Tensor,
+) -> tf.Tensor:
+    expected_rank = len(keras.ops.shape(edge))
+    current_rank = len(keras.ops.shape(edge_weight))
+    for _ in range(expected_rank - current_rank):
+        edge_weight = keras.ops.expand_dims(edge_weight, axis=-1)
+    return edge * edge_weight
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def segment_mean(
+    data: tf.Tensor,
+    segment_ids: tf.Tensor,
+    num_segments: int | None = None,
+    sorted: bool = False,
+) -> tf.Tensor:
+    if num_segments is None:
+        num_segments = keras.ops.cond(
+            keras.ops.greater(keras.ops.shape(segment_ids)[0], 0),
+            lambda: keras.ops.max(segment_ids) + 1,
+            lambda: 0
+        )
+    if backend.backend() == 'tensorflow':
+        segment_mean_fn = (
+            tf.math.unsorted_segment_mean if not sorted else
+            tf.math.segment_mean
+        )
+        return segment_mean_fn(
+            data=data,
+            segment_ids=segment_ids,
+            num_segments=num_segments
+        )
+    x = keras.ops.segment_sum(
+        data=data, 
+        segment_ids=segment_ids, 
+        num_segments=num_segments,
+        sorted=sorted
+    )
+    sizes = keras.ops.cast(
+        keras.ops.bincount(segment_ids, minlength=num_segments), 
+        dtype=x.dtype
+    )
+    return x / sizes[:, None]
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def gaussian(
+    x: tf.Tensor, 
+    mean: tf.Tensor, 
+    std: tf.Tensor
+) -> tf.Tensor:
+    expected_rank = len(keras.ops.shape(x))
+    current_rank = len(keras.ops.shape(mean))
+    for _ in range(expected_rank - current_rank):
+        mean = keras.ops.expand_dims(mean, axis=0)
+        std = keras.ops.expand_dims(std, axis=0)
+    a = (2 * np.pi) ** 0.5
+    return keras.ops.exp(-0.5 * (((x - mean) / std) ** 2)) / (a * std)
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def euclidean_distance(
+    x1: tf.Tensor, 
+    x2: tf.Tensor, 
+    axis: int = -1,
+    keepdims: bool = True,
+) -> tf.Tensor:
+    relative_distance = keras.ops.subtract(x1, x2)
+    return keras.ops.sqrt(
+        keras.ops.sum(
+            keras.ops.square(relative_distance), 
+            axis=axis, 
+            keepdims=keepdims
+        )
+    )
+
+@keras.saving.register_keras_serializable(package='molcraft')
+def displacement(
+    x1: tf.Tensor,
+    x2: tf.Tensor,
+    normalize: bool = False,
+    axis: int = -1,
+    keepdims: bool = True,
+) -> tf.Tensor:
+    displacement = keras.ops.subtract(x1, x2)
+    if not normalize:
+        return displacement
+    return displacement / euclidean_distance(x1, x2, axis=axis, keepdims=keepdims)

molcraft/records.py

@@ -0,0 +1,187 @@
+import warnings
+import os
+import math
+import glob
+import time
+import typing 
+import tensorflow as tf
+import numpy as np
+import pandas as pd
+import multiprocessing as mp
+
+from molcraft import tensors
+
+if typing.TYPE_CHECKING:
+    from molcraft import featurizers
+
+
+def write(
+    inputs: list[str | tuple], 
+    featurizer: 'featurizers.GraphFeaturizer',
+    path: str, 
+    exist_ok: bool = False,
+    overwrite: bool = False, 
+    num_files: typing.Optional[int] = None, 
+    num_processes: typing.Optional[int] = None,
+    multiprocessing: bool = False,
+    device: str = '/cpu:0'
+) -> None:
+    
+    if os.path.isdir(path):
+        if not exist_ok:
+            raise FileExistsError(f'Records already exist: {path}')
+        if not overwrite:
+            return
+        else:
+            _remove_files(path)
+    else:
+        os.makedirs(path)
+
+    with tf.device(device):
+        
+        if isinstance(inputs, (pd.DataFrame, pd.Series)):
+            inputs = list(inputs.iterrows())
+
+        example = featurizer._call(inputs[0])
+        save_spec(os.path.join(path, 'spec.pb'), example.spec)
+
+        if num_processes is None:
+            num_processes = mp.cpu_count()
+
+        if num_files is None:
+            num_files = min(len(inputs), max(1, math.ceil(len(inputs) / 1_000)))
+            
+        num_examples = len(inputs)
+        chunk_sizes = [0] * num_files
+        for i in range(num_examples):
+            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
+        
+        paths = [
+            os.path.join(path, f'tfrecord-{i:06d}.tfrecord')
+            for i in range(num_files)
+        ]
+        
+        if not multiprocessing:
+            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, start_index in zip(paths, input_chunks, start_indices):
+        
+            while len(processes) >= num_processes:
+                for process in processes:
+                    if not process.is_alive():
+                        processes.remove(process)
+                else:
+                    time.sleep(0.1)
+                    continue
+                    
+            process = mp.Process(
+                target=_write_tfrecord,
+                args=(input_chunk, path, featurizer, start_index)
+            )
+            processes.append(process)
+            process.start()
+
+        for process in processes:
+            process.join()         
+    
+def read(
+    path: str, 
+    shuffle_files: bool = False
+) -> tf.data.Dataset:
+    spec = load_spec(os.path.join(path, 'spec.pb'))
+    filenames = sorted(glob.glob(os.path.join(path, '*.tfrecord')))
+    num_files = len(filenames)
+    ds = tf.data.Dataset.from_tensor_slices(filenames)
+    if shuffle_files:
+        ds = ds.shuffle(num_files)
+    ds = ds.interleave(
+        tf.data.TFRecordDataset, num_parallel_calls=1)
+    ds = ds.map(
+        lambda x: _parse_example(x, spec),
+        num_parallel_calls=tf.data.AUTOTUNE)
+    if not tensors.is_scalar(spec):
+        ds = ds.unbatch()
+    return ds
+
+def save_spec(path: str, spec: tensors.GraphTensor.Spec) -> None:
+    proto = spec.experimental_as_proto()
+    with open(path, 'wb') as fh:
+        fh.write(proto.SerializeToString())
+
+def load_spec(path: str) -> tensors.GraphTensor.Spec:
+    with open(path, 'rb') as fh:
+        serialized_proto = fh.read()
+    spec = tensors.GraphTensor.Spec.experimental_from_proto(
+        tensors.GraphTensor.Spec
+        .experimental_type_proto()
+        .FromString(serialized_proto)
+    )
+    return spec
+    
+def _write_tfrecord(
+    inputs: list, 
+    path: str,
+    featurizer: 'featurizers.GraphFeaturizer', 
+    start_index: int,
+) -> None:
+    with tf.io.TFRecordWriter(path) as writer:
+        for i, x in enumerate(inputs):
+            try:
+                tensor = featurizer._call(x)
+                serialized = _serialize_example(tensor)
+                writer.write(serialized)
+            except Exception as e:
+                index = getattr(x, 'Index', (i + start_index))
+                warnings.warn(
+                    f'Could not write record for index {index}, '
+                    f'proceeding without it. Exception raised:\n{e}'
+                )
+
+def _serialize_example(tensor):
+    flat_values = tf.nest.flatten(tensor, expand_composites=True)
+    flat_values = [
+        tf.io.serialize_tensor(value).numpy() for value in flat_values
+    ]
+    feature = tf.train.Feature(
+        bytes_list=tf.train.BytesList(value=flat_values)
+    )
+    example_proto = tf.train.Example(
+        features=tf.train.Features(feature={'feature': feature})
+    )
+    return example_proto.SerializeToString()
+    
+def _parse_example(
+    x: tf.Tensor, 
+    spec: tensors.GraphTensor.Spec
+) -> 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.nest.pack_sequence_as(spec, tf.nest.flatten(out), expand_composites=True)
+    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)