nucleobench 0.1.0__py3-none-any.whl

nucleobench/common/argparse_lib.py

@@ -0,0 +1,64 @@
+"""Utilities for parsing arguments."""
+
+from typing import Iterable
+
+import argparse
+import dataclasses
+
+@dataclasses.dataclass
+class ParsedArgs:
+    main_args: argparse.Namespace
+    model_init_args: argparse.Namespace
+    opt_init_args: argparse.Namespace
+    opt_run_args: argparse.Namespace
+    
+    
+def parse_long_start_sequence(known_args: argparse.Namespace) -> argparse.Namespace:
+    """Parse a long start sequence from a file."""
+    assert known_args.seed_sequence.startswith('local://')
+    local_fileloc = known_args.seed_sequence[len('local://'):]
+    with open(local_fileloc, 'r') as f:
+        known_args.seed_sequence = f.read()
+    return known_args
+
+
+def possibly_parse_positions_to_mutate(known_args: argparse.Namespace) -> argparse.Namespace:
+    """Possibly parse `positions_to_mutate` from a file, or leave it untouched, depending on the value."""
+    if isinstance(known_args.positions_to_mutate, str) and known_args.positions_to_mutate.startswith('local://'):
+        local_fileloc = known_args.positions_to_mutate[len('local://'):]
+        with open(local_fileloc, 'r') as f:
+            loc_str = f.read()
+        known_args.positions_to_mutate = [int(x) for x in loc_str.split('\n')]
+    elif known_args.positions_to_mutate is None or known_args.positions_to_mutate == '' or known_args.positions_to_mutate == []:
+        known_args.positions_to_mutate = None
+    else:
+        assert isinstance(known_args.positions_to_mutate, list), (type(known_args.positions_to_mutate), known_args.positions_to_mutate)
+        known_args.positions_to_mutate = [int(x) for x in known_args.positions_to_mutate.split(',')]
+    return known_args
+
+
+def handle_leftover_args(known_args: argparse.Namespace, leftover_args: Iterable):
+    """Handle leftover arguments, either by failing or by ignoring them."""
+    if known_args.ignore_empty_cmd_args:
+        # Check that every "value" is either `None` or `empty`. If so, allow it to continue.
+        for i in leftover_args:
+            if i.startswith('--'):
+                if '=' in i:
+                    arg_val = i.split('=')[1]
+                    if arg_val not in [None, '']:
+                        raise ValueError(f'Unused arg, not empty: {leftover_args}')
+                continue
+            else: 
+                if i not in [None, '']:
+                    raise ValueError(f'Unused arg, not empty: {leftover_args}')
+    else:
+        raise ValueError(f'Unused args: {leftover_args}')
+    
+
+def str_to_bool(s):
+    if s.lower() in ('yes', 'true', 't', '1'):
+        return True
+    elif s.lower() in ('no', 'false', 'f', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')

nucleobench/common/attribution_lib_torch.py

@@ -0,0 +1,199 @@
+"""Library for smoothgrad and genome-specific attribution methods.
+
+Ref:
+1. [Correcting gradient-based interpretations of deep neural networks for genomics](https://genomebiology.biomedcentral.com/articles/10.1186/s13059-023-02956-3)
+2. [SmoothGrad: removing noise by adding noise](https://arxiv.org/abs/1706.03825)
+3. [Quick and effective approximation of in silico saturation mutagenesis experiments with first-order taylor expansion](https://pubmed.ncbi.nlm.nih.gov/39286491/)
+
+TODO(joelshor): Consider using public version of attribution tools, such as:
+- PyTorch Smoothgrad: https://github.com/pkmr06/pytorch-smoothgrad
+- PyTorch GradCam and others: https://github.com/jacobgil/pytorch-grad-cam?tab=readme-ov-file
+- PyTorch Smoothgrad and others: https://tf-explain.readthedocs.io/en/latest/
+
+To test locally:
+```zsh
+python -m nucleobench.common.attribution_lib
+```
+"""
+
+import gc
+import numpy as np
+import torch
+from typing import Callable, Optional
+
+
+TISMOutputType = list[dict[str, float]]
+SmoothgradVocabType = list[dict[str, torch.Tensor]]
+TISMLocationsType = list[int]
+
+
+def noise_inputs(
+    input_tensor: torch.Tensor,
+    noise_stdev: float,
+    times: int,
+    ) -> torch.Tensor:
+    """Generates noisy inputs.
+
+    NOTE: For simplicity, for now, we work with SINGLE TENSORS. Assume no batch dimension.
+
+    Args:
+        input_tensor: Input tensor. Doesn't have to be genomic. Should NOT be batched.
+        noise_stdev: Noise to add.
+        times: Number of times to add noise.
+    """
+    if noise_stdev < 0:
+        raise ValueError(f'Requires non-negative noise stdev: {noise_stdev}')
+    x = input_tensor  # Syntactic sugar.
+    
+    # Stack N versions of the input, to add uncorrelated noise to.
+    with torch.no_grad():
+        x = x.unsqueeze(0)
+        x = x.repeat([times] + [1] * (x.ndim-1))
+
+        # Add noise for hte smoothgrad algorithm.
+        if noise_stdev > 0:
+            noise_to_add = torch.normal(mean=torch.zeros(x.shape), std=noise_stdev)
+            x += noise_to_add
+    return x
+
+def noisy_grads_torch(
+        input_tensor: torch.Tensor, 
+        model: Callable[[torch.Tensor], torch.Tensor], 
+        noise_stdev: float, 
+        times: int,
+        idxs: Optional[TISMLocationsType] = None,
+        ) -> torch.Tensor:
+    """Generates noisy gradients from a function.
+
+    NOTE: For simplicity, for now, we work with SINGLE TENSORS. Assume no batch dimension.
+
+    This replicates the input `times` times, and runs it through the network all at once.
+
+    TODO(joelshor): Add batching, for the situation where `times` is larger than the possible batch size
+        of a single inference with a network.
+    TODO(joelshor): Add ability to efficiently compute multiple inputs at once.
+
+    Args:
+        input_tensor (torch.Tensor): Input tensor. Doesn't have to be genomic. Should NOT be batched.
+        model: PyTorch model to use. The model must return a scalar per batch element.
+        noise_stdev: Noise to add.
+        times: Number of times to add noise.
+        idx: If present, only backprop through this location.
+    """
+    x = noise_inputs(
+        input_tensor=input_tensor,
+        noise_stdev=noise_stdev,
+        times=times)
+
+    # Run inference to get grads.
+    if idxs is None:
+        x_grad = x
+        x_grad.requires_grad = True
+    else:
+        x, x_grad = apply_gradient_mask(x, idxs)
+        
+    y = model(x)
+    y_sum = y.sum()
+    y_sum.backward(retain_graph=False)
+    noisy_grads = x_grad.grad.numpy()
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    assert noisy_grads.shape == x_grad.shape
+    return noisy_grads
+
+
+def smoothgrad_torch(
+    input_tensor: torch.Tensor, 
+    model: Callable[[torch.Tensor], torch.Tensor], 
+    noise_stdev: float, 
+    times: int,
+    idxs: Optional[TISMLocationsType] = None,
+    ) -> torch.Tensor:
+    """Custom implementation of SmoothGrad.
+    https://arxiv.org/pdf/1706.03825
+
+    NOTE: For simplicity, for now, we work with SINGLE TENSORS. Assume no batch dimension.
+
+    This replicates the input `times` times, and runs it through the network all at once.
+
+    TODO(joelshor): Add batching, for the situation where `times` is larger than the possible batch size
+        of a single inference with a network.
+    TODO(joelshor): Add ability to efficiently compute multiple inputs at once.
+
+    Args:
+        input_tensor (torch.Tensor): Input tensor. Doesn't have to be genomic. Should NOT be batched.
+        model: PyTorch model to use. The model must return a scalar per batch element.
+        noise_stdev: Noise to add.
+        times: Number of times to add noise.
+        idx: If present, only backprop to this location.
+        
+    Returns:
+        Per-nucleotide smoothgrad.
+    """
+    noisy_grads = noisy_grads_torch(
+        input_tensor=input_tensor,
+        model=model,
+        noise_stdev=noise_stdev,
+        times=times,
+        idxs=idxs,
+    )
+
+    return np.mean(noisy_grads, axis=0)
+
+
+# TODO(joelshor): Add `attribution_lib.py` test, taken from `malinois/model_def_test.py`.
+def smoothgrad_tensor_to_dict(smooth_grad: torch.Tensor, vocab: list[str]) -> SmoothgradVocabType:
+    """Map the smoothgrad indices to the vocab."""
+    assert smooth_grad.ndim == 2
+    assert list(smooth_grad.shape)[0] == len(vocab)
+    def _to_dict(x: torch.Tensor) -> dict[str, torch.Tensor]:
+        return {vocab[i]: x[i] for i in range(len(vocab))}
+    return [_to_dict(x) for x in smooth_grad.T]
+
+
+def smoothgrad_to_tism(sg: SmoothgradVocabType, base_seq: str) -> TISMOutputType:
+    """Returns result according to Taylor in-silico mutagenesis.
+    
+    Paper: https://www.cell.com/iscience/fulltext/S2589-0042(24)02032-7"""
+    assert len(sg) == len(base_seq)
+
+    tism = []
+    for base_nt, sg_dict in zip(base_seq, sg):
+        cur_tism = {}
+        for nt, sg in sg_dict.items():
+            if nt == base_nt: continue
+            cur_tism[nt] = float(sg - sg_dict[base_nt])
+        tism.append(cur_tism)
+
+    return tism
+
+
+def apply_gradient_mask(x: torch.Tensor, idxs: TISMLocationsType) -> tuple[torch.Tensor, torch.Tensor]:
+    """Applies a gradient mask to the input tensor.
+    
+    NOTE: Do NOT just multiply by 0. This will run out of memory in large models.
+
+    Returns:
+        Tuple of (x, masked_x), where masked_x is the input tensor with the gradient mask applied.
+    """
+    assert min(idxs) >= 0
+    assert max(idxs) < x.shape[2]
+    assert x.ndim == 3, x.shape
+    
+    no_gradient = x.clone().detach()
+    no_gradient.requires_grad = False
+    
+    x_grad = x[:, :, idxs].clone().detach()
+    x_grad.requires_grad = True
+    x_grad_i = {idx: i for i, idx in enumerate(idxs)}
+    
+    # Instead of using `torch.where`, we use this method to make our gradient tensor
+    # as small as possible, to preserve memory.
+    tensor_slices = [x_grad[:, :, x_grad_i[i]:x_grad_i[i]+1] if i in idxs
+                     else no_gradient[:, :, i:i+1]
+                     for i in range(no_gradient.shape[2])]
+    x = torch.concat(tensor_slices, dim=2)
+    
+    return x, x_grad

nucleobench/common/attribution_lib_torch_test.py

@@ -0,0 +1,175 @@
+"""Test attribution_lib.py.
+
+To test:
+pytest nucleobench/common/attribution_lib_torch_test.py
+"""
+
+import pytest
+import numpy as np
+import torch
+from torch import nn
+
+from nucleobench.common import attribution_lib_torch
+from nucleobench.common import testing_utils
+
+class TestNeuralNetwork(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(10, 3),
+        )
+
+    def forward(self, x):
+        x = self.flatten(x)
+        logits = self.linear_relu_stack(x)
+        return logits
+    
+
+class Nonlinear(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(10, 2),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        x = self.flatten(x)
+        logits = self.linear_relu_stack(x)
+        return logits
+
+
+class FixedNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layer_array = np.array([1.0, 2.0])
+        self.linear_layer = torch.tensor([self.layer_array], dtype=torch.float).t()
+
+    def forward(self, x):
+        logits = torch.matmul(x, self.linear_layer)
+        return logits
+
+@pytest.mark.parametrize('times', [1, 3, 5])
+def test_expected_gradient(times):
+    """Regardless of number of times, noisy grads should be the same for a linear model."""
+    input_tensor = torch.randn(2)
+    model = FixedNN()
+
+    grads = attribution_lib_torch.noisy_grads_torch(
+        input_tensor=input_tensor,
+        model=model,
+        noise_stdev=1.0,
+        times=times)
+    
+    assert grads.shape == (times, 2)
+    # For linear models, grads should all be the linear layer.
+    for i in range(times):
+        assert np.array_equal(grads[i, ...], model.layer_array), (
+            grads[i, ...], model.layer_array, grads)
+        
+def test_callable():
+    """Check that additional arguments are used."""
+    input_tensor = torch.randn(2)
+    model = FixedNN()
+
+    def override_callable(x):
+        return 2.0 * model(x)
+    grads = attribution_lib_torch.noisy_grads_torch(
+        input_tensor=input_tensor,
+        model=override_callable,
+        noise_stdev=1.0,
+        times=3,
+    )
+    
+    assert grads.shape == (3, 2)
+    # For linear models, grads should all be the linear layer.
+    for i in range(3):
+        assert np.array_equal(grads[i, ...], [2.0, 4.0]), (
+            grads[i, ...], [2.0, 4.0], grads)
+
+def test_noisy_grads():
+    input_tensor = torch.randn(10)
+    model = TestNeuralNetwork()
+    
+    noisy_grads = attribution_lib_torch.noisy_grads_torch(
+        input_tensor=input_tensor,
+        model=model,
+        noise_stdev=0.25,
+        times=5)
+    assert noisy_grads.shape == (5, 10)
+    # For linear models, grads should all be the same.
+    for i in range(5):
+        assert np.array_equal(noisy_grads[i, ...], noisy_grads[0, ...])
+
+
+def test_smoothgrad_torch():
+    input_tensor = torch.randn(10)
+    model = TestNeuralNetwork()
+    
+    grad = attribution_lib_torch.smoothgrad_torch(
+        input_tensor=input_tensor,
+        model=model,
+        noise_stdev=0.25,
+        times=5)
+
+    assert grad.shape == (10,)
+
+
+def test_noisy_grads_different():
+    """In a nonlinear network, gradients should be different."""
+    input_tensor = torch.randn(10)
+    model = Nonlinear()
+    
+    noisy_grads = attribution_lib_torch.noisy_grads_torch(
+        input_tensor=input_tensor,
+        model=model,
+        noise_stdev=1.0,
+        times=5)
+    assert noisy_grads.shape == (5, 10)
+    # For linear models, grads should all be the same.
+    for i in range(1, 5):
+        assert not np.array_equal(noisy_grads[i, ...], noisy_grads[0, ...])
+        
+# TODO(joelshor): Add a unit test for nucleotide-specific TISM.
+def test_smoothgrad_torch_idx_sanity():
+    input_tensor = torch.randn(4, 5)
+    model = testing_utils.CountLetterModel()
+    
+    grad = attribution_lib_torch.smoothgrad_torch(
+        input_tensor=input_tensor,
+        model=model.inference_on_tensor,
+        noise_stdev=0.0,
+        times=1)
+    assert grad.shape == (4, 5)
+    
+    for idx in range(5):
+        grad_singlebp = attribution_lib_torch.smoothgrad_torch(
+            input_tensor=input_tensor,
+            model=model.inference_on_tensor,
+            noise_stdev=0.25,
+            times=5,
+            idxs=[idx])
+        assert grad_singlebp.shape == (4, 1)
+        assert np.all(grad[:, idx] == grad_singlebp[:, 0])
+        
+        
+def test_apply_gradient_mask():
+    """[idx] and idx should give the same result."""
+    input_tensor = torch.randn(1, 10)
+    model = TestNeuralNetwork()
+    output_tensor = model(input_tensor)
+    output_tensor = output_tensor.reshape(1, 1, 3)
+    
+    x1, x_grad = attribution_lib_torch.apply_gradient_mask(
+        output_tensor, [0])
+    assert x1.shape == (1, 1, 3)
+    assert x_grad.shape == (1, 1, 1)
+    
+    x2, x_grad = attribution_lib_torch.apply_gradient_mask(
+        output_tensor, [0, 1])
+    assert x2.shape == (1, 1, 3)
+    assert x_grad.shape == (1, 1, 2)
+    
+    assert (x1 == x2).all(), x1 == x2

nucleobench/common/constants.py

@@ -0,0 +1,7 @@
+VOCAB = ['A','C','G','T']
+SERVICE_KEY_FILE_LOCATION = ''
+GCP_OUTPUT_BUCKET_NAME = ''
+
+
+MPRA_UPSTREAM  = 'ACGAAAATGTTGGATGCTCATACTCGTCCTTTTTCAATATTATTGAAGCATTTATCAGGGTTACTAGTACGTCTCTCAAGGATAAGTAAGTAATATTAAGGTACGGGAGGTATTGGACAGGCCGCAATAAAATATCTTTATTTTCATTACATCTGTGTGTTGGTTTTTTGTGTGAATCGATAGTACTAACATACGCTCTCCATCAAAACAAAACGAAACAAAACAAACTAGCAAAATAGGCTGTCCCCAGTGCAAGTGCAGGTGCCAGAACATTTCTCTGGCCTAACTGGCCGCTTGACG'
+MPRA_DOWNSTREAM= 'CACTGCGGCTCCTGCGATCTAACTGGCCGGTACCTGAGCTCGCTAGCCTCGAGGATATCAAGATCTGGCCTCGGCGGCCAAGCTTAGACACTAGAGGGTATATAATGGAAGCTCGACTTCCAGCTTGGCAATCCGGTACTGTTGGTAAAGCCACCATGGTGAGCAAGGGCGAGGAGCTGTTCACCGGGGTGGTGCCCATCCTGGTCGAGCTGGACGGCGACGTAAACGGCCACAAGTTCAGCGTGTCCGGCGAGGGCGAGGGCGATGCCACCTACGGCAAGCTGACCCTGAAGTTCATCT'

nucleobench/common/gcp_utils.py

@@ -0,0 +1,150 @@
+"""Utils for reading from and writing to GCP.
+
+To test this with the cloud:
+```zsh
+python -m nucleobench.common.gcp_utils
+```
+"""
+
+from typing import Any, Generator
+
+import argparse
+import copy
+import os
+import pickle
+import subprocess
+import torch
+import time
+from google.cloud import storage
+
+from nucleobench.common import constants
+
+
+def get_filepath(
+    base_dir: str,
+    opt_method: str,
+    model: str,
+    exp_start_time: str,
+    timestamp: str,
+    ) -> str:
+    return os.path.join(base_dir, f'{opt_method}_{model}', exp_start_time, f'{timestamp}.pkl')
+
+def save_proposals(
+        write_dicts: list[dict],
+        args: argparse.Namespace, 
+        output_path: str,
+        ):
+    """
+    Save the proposals and associated arguments to a file.
+
+    This function saves the generated proposals and the input arguments used for their generation
+    to a file. The file is named based on the current timestamp and a random tag. The saved file can
+    be placed in a local directory or on Google Cloud Storage if a 'gs://' path is provided in the arguments.
+
+    Args:
+        write_dicts: A list of dictiomary of things to write.
+        args: Args for the job. Best to keep them close to the output.
+        output_path: Directory to write the output to, either locally or on GCP. Format of output is:
+            {output_path}/{opt_method}_{model}/{exp start time}/{reults of exp N}.pkl
+
+    Returns:
+        None
+    """
+    def _tensor2np(obj):
+        return obj.detach().clone().numpy() if isinstance(obj, torch.Tensor) else obj
+    args = {k: _tensor2np(v) for k, v in vars(args).items()}
+    write_dicts = [{k: _tensor2np(v) for k, v in x.items()} for x in write_dicts]
+    save_dicts = copy.deepcopy(write_dicts)
+
+    filename = get_filepath(
+        base_dir=output_path,
+        opt_method=args['optimization'],
+        model=args['model'],
+        exp_start_time=save_dicts[0]['exp_starttime_str'],
+        timestamp=time.strftime("%Y%m%d_%H%M%S"),
+    )
+    
+    # TODO(joelshor): Used to save as torch tensors. Now, saving as a pickled dictionary
+    # of numpy arrays. Eventually, save as PyArrow arrays to parquet.
+    if filename.startswith('gs://'):
+        bucket_name = filename.split('/')[2]
+        write_str_to_gcp(
+            gcs_output_path=filename,
+            content=pickle.dumps(save_dicts),
+            binary=True,
+            bucket_name=bucket_name,
+        )
+    else:
+        if os.path.dirname(filename) != '' and os.path.dirname(filename) != '.':
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+        with open(filename, 'wb') as f:
+            pickle.dump(save_dicts, f)
+    
+    print(f'Proposals deposited at:\n\t{filename}')
+
+
+def get_role_client(service_json_path: str = constants.SERVICE_KEY_FILE_LOCATION):
+    try:
+        gcp_client = storage.Client.from_service_account_json(service_json_path)
+    except FileNotFoundError as e:
+        raise ValueError(
+            'GCP service key not found. You are probably trying to write output to a private '
+            'bucket. To write to your own bucket, change the bucket name and service key '
+            f'credentials in `nucleobench/common/constants.py`: {service_json_path}') from e
+    return gcp_client
+
+def write_str_to_gcp(
+        gcs_output_path: str,
+        content: Any,
+        binary: bool,
+        bucket_name: str = constants.GCP_OUTPUT_BUCKET_NAME,
+        ):
+    assert gcs_output_path.startswith('gs://'), 'gcs_output_path must be a GCS path.'
+    gcs_output_path = gcs_output_path[len('gs://'):]
+    bucket_name, blob_fn = gcs_output_path.split('/', 1)
+
+    # Instantiates a client.
+    storage_client = get_role_client()
+    bucket = storage_client.bucket(bucket_name)
+    blob = bucket.blob(blob_fn)
+
+    write_type = 'wb' if binary else 'w'
+    with blob.open(write_type) as f:
+        f.write(content)
+        
+        
+def download_gcp_folder_to_local(dir: str, local_dir: str, bucket: str = constants.GCP_OUTPUT_BUCKET_NAME):
+    subprocess.call(['mkdir', '-p', local_dir])
+    subprocess.call(['gcloud', 'storage', 'cp', '-r', f'gs://{bucket}/{dir}', local_dir])
+
+
+def list_files_recursively(local_dir: str) -> Generator[str, None, None]:
+    """Recursively lists all files in a given directory."""
+    for root, dirs, files in os.walk(local_dir):
+        for file in files:
+            yield os.path.join(root, file)
+
+
+if __name__ == '__main__':
+    args = {
+        'fake_arg': 'fake_value',
+        'fake_arg2': 'fake_value2',
+        'optimization': 'dummy',
+        'model': 'dummy_model',
+    }
+    proposal1 = {
+        'string': 'ACTC',
+        'energy': 1.0,
+    }
+    proposal2 = {
+        'string': 'ACTC',
+        'energy': 2.0,
+    }
+
+    write_dict = {
+        'proposals': [proposal1, proposal2],
+        'exp_starttime_str': '102010',
+    }
+
+    save_proposals([write_dict, write_dict], argparse.Namespace(**args), './gcp_utils_test/dummy')
+    save_proposals([write_dict, write_dict], argparse.Namespace(**args), f'gs://{constants.GCP_OUTPUT_BUCKET_NAME}/gcp_utils_test/dummy')

nucleobench/common/memory_utils.py

@@ -0,0 +1,10 @@
+"""Utilities for managing memory."""
+
+import gc
+import torch
+
+def free_memory(debug: bool = False):
+    collected = gc.collect()
+    if debug:
+        print(f'[free_memory] Collected: {collected}')
+    torch.cuda.empty_cache()