yunta 0.0.1__py3-none-any.whl

yunta/cli.py

@@ -0,0 +1,252 @@
+"""Command-line interface for sppid."""
+
+__version__ = '0.0.1'
+
+from typing import Any, Mapping, Optional, Tuple, Union
+
+from argparse import ArgumentParser, FileType, Namespace
+from io import TextIOWrapper
+import os
+import sys
+
+from carabiner import print_err
+from carabiner.cast import cast, flatten
+from carabiner.cliutils import clicommand, CLIOption, CLICommand, CLIApp
+
+from .io import write_metrics
+from .plots import plot_matrix
+from .screening import (
+    dca_one_vs_many, 
+    dca_many_vs_many, 
+    model_one_vs_many, 
+    model_many_vs_many, 
+    rf2track_one_vs_many,
+)
+
+def _load_msa_list(*args):
+    args = [a[0] if isinstance(a, list) else a for a in args]
+    return [flatten([line.strip() for line in msa]) for msa in args]
+
+def _plot_results(results, result_interaction, metric, 
+                  output_dir: str = '.', *args, **kwargs) -> None:
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    filename_prefix = os.path.join(output_dir, metric.ID)
+    if hasattr(metric, 'apc'):
+        apc = metric.apc
+        filename_prefix += f".{apc=}"    
+    plot_matrix(results, 
+                filename_prefix=filename_prefix, 
+                hline=metric.chain_a_len, 
+                vline=metric.chain_a_len)
+    plot_matrix(result_interaction, 
+                filename_prefix=f"{filename_prefix}.interaction",
+                ylabel=metric.uniprot_id_1, 
+                xlabel=metric.uniprot_id_2)
+    return None
+
+
+@clicommand(message="Making RosettaFold-2track prediction with the following parameters")
+def _rf2t_single(args: Namespace) -> None:
+
+    if args.list_file:
+        msa1, msa2 = _load_msa_list(args.msa1, args.msa2)
+    else:
+        msa1, msa2 = args.msa1, args.msa2
+
+    print_err(f"Running RF-2t using {msa1} as reference.")
+    outputs = rf2track_one_vs_many(
+        msa_file1=msa1,
+        msa_file2=msa2,
+        cpu=args.cpu,
+    )
+    metrics = [_output[-1] for _output in outputs]
+    write_metrics(metrics, 
+                  filename=args.output)
+    if args.plot is not None:
+        for _output in outputs:
+            _plot_results(*_output, output_dir=args.plot)
+
+    return None
+
+
+
+@clicommand(message="Calculating DCA for a pair of MSAs with the following parameters")
+def _dca_single(args: Namespace) -> None:
+
+    if args.list_file:
+        msa1, msa2 = _load_msa_list(args.msa1, args.msa2)
+    else:
+        msa1, msa2 = args.msa1, args.msa2
+
+    outputs = dca_one_vs_many(
+        msa_file1=msa1,
+        msa_file2=msa2,
+        apc=args.apc,
+    )
+    metrics = [_output[-1] for _output in outputs]
+    write_metrics(metrics, 
+                  filename=args.output)
+    if args.plot is not None:
+        for _output in outputs:
+            _plot_results(*_output, output_dir=args.plot)
+
+    return None
+
+
+@clicommand(message="Calculating DCA between pairs of MSAs with the following parameters")
+def _dca_many_vs_many(args: Namespace) -> None:
+
+    if args.list_file:
+        msa1, msa2 = _load_msa_list(args.msa1, args.msa2)
+    else:
+        msa1, msa2 = args.msa1, args.msa2
+
+    outputs = dca_many_vs_many(
+        msa_files1=msa1,
+        msa_files2=msa2,
+        apc=args.apc,
+    )
+
+    metrics = [_output[-1] for _output in outputs]
+    write_metrics(metrics, 
+                  filename=args.output)
+    if args.plot is not None:
+        for _output in outputs:
+            _plot_results(*_output, output_dir=args.plot)
+
+    return None
+
+
+@clicommand(message="Modelling one PPI with the following parameters")
+def _af2_single(args: Namespace) -> None:
+
+    if args.list_file:
+        msa1, msa2 = _load_msa_list(args.msa1, args.msa2)
+    else:
+        msa1, msa2 = args.msa1, args.msa2
+
+    metric = model_one_vs_many(
+        msa_file1=msa1,
+        msa_file2=msa2,
+        max_recycles=args.recycles,
+        output_dir=args.output,
+        param_dir=args.params,
+    )
+
+    output_filename = os.path.join(args.output_dir, f"{metric.ID}_metrics.csv")
+    print_err(f"Saving metrics as {output_filename}")
+    write_metrics(metric, 
+                  filename=output_filename)
+
+    return None
+
+
+@clicommand(message="Modelling sets of PPIs with the following parameters")
+def _af2_many_vs_many(args: Namespace) -> None:
+
+    if args.list_file:
+        msa1, msa2 = _load_msa_list(args.msa1, args.msa2)
+    else:
+        msa1, msa2 = args.msa1, args.msa2
+
+    metrics = model_many_vs_many(
+        msa_files1=msa1,
+        msa_files2=msa2,
+        output_dir=args.output,
+        max_recycles=args.recycles,
+        param_dir=args.params,
+    )
+
+    output_filename = os.path.join(args.output, "_all_metrics.csv")
+    print_err(f"Saving metrics as {output_filename}")
+    write_metrics(metrics, 
+                  filename=output_filename)
+
+    return None
+
+
+def main() -> None:
+    inputs = CLIOption('msa1', 
+                       default=sys.stdin,
+                       type=FileType('r'), 
+                       nargs='?',
+                       help='MSA file. Default: STDIN.')
+    input2 = CLIOption('--msa2', '-2', 
+                       type=FileType('r'), 
+                       default=None,
+                       help='Second MSA file.')
+    inputs_list = CLIOption('msa1', 
+                            default=sys.stdin,
+                            type=FileType('r'), 
+                            nargs='*',
+                            help='MSA file(s).')
+    inputs_list2 = CLIOption('--msa2', '-2', 
+                        type=FileType('r'), 
+                        default=None,
+                        nargs='*',
+                        help='Second MSA file(s). Default: if not provided, all pairwise from msa1.')
+    list_file = CLIOption('--list-file', '-l', 
+                          action='store_true',
+                          help='Treat inputs as plain-text list of MSA files, rather than MSA filenames. '
+                               'Default: treat as MSA filenames.')
+    output = CLIOption('--output', '-o', 
+                       type=str,
+                       required=True,
+                       help='Output directory.')
+    plot = CLIOption('--plot', '-p', 
+                     type=str,
+                     default=None,
+                     help='Directory for saving plots. Default: don\'t plot.')
+    cpu = CLIOption('--cpu', '-c', 
+                    action='store_true',
+                    help='Whether to use CPU only. Default: use GPU.')
+    output_file = CLIOption('--output', '-o', 
+                            default=sys.stdout,
+                            type=FileType('w'), 
+                            nargs='?',
+                            help='Output filename. Default: STDOUT.')
+    apc = CLIOption('--apc', '-a', 
+                    action='store_true',
+                    help='Whether to use APC correction in DCA. Default: don\'t apply correction.')
+    params = CLIOption('--params', '-w', 
+                       type=str,
+                       default=None,
+                       help='Path to AlphaFold2 params file (.npz).')
+    recycles = CLIOption('--recycles', '-x', 
+                       type=int,
+                       default=10,
+                       help='Maximum number of recyles through the model.')
+
+    rf2t_single = CLICommand('rf2t-single', 
+                            description='Calculate RF-2track contacts for between one protein and a series of others.',
+                            main=_rf2t_single,
+                            options=[inputs, inputs_list2, list_file, output_file, plot, cpu])
+    dca_single = CLICommand('dca-single', 
+                            description='Calculate DCA for one protein-protein interaction.',
+                            main=_dca_single,
+                            options=[inputs, inputs_list2, list_file, output_file, plot, apc])
+    dca_many = CLICommand('dca-many', 
+                          description='Calculate DCA between two sets of proteins, or all pairs in one set of proteins.',
+                          main=_dca_many_vs_many,
+                          options=[inputs_list, inputs_list2, list_file, apc, output_file, plot])
+    af2_single = CLICommand('af2-single', 
+                            description='Model one protein-protein interaction.',
+                            main=_af2_single,
+                            options=[inputs, inputs_list2, list_file, output, params, recycles, plot])
+    af2_many = CLICommand('af2-many', 
+                          description='Model all interactions between two sets of proteins, or all pairs in one set of proteins.',
+                          main=_af2_many_vs_many,
+                          options=[inputs_list, inputs_list2, list_file, output, params, recycles, plot])
+
+    app = CLIApp("sppid",
+                 version=__version__,
+                 description="Screening protein-protein interactions using DCA, RosettaFold-2track, and AlphaFold2.",
+                 commands=[dca_single, dca_many, rf2t_single, af2_single, af2_many])
+
+    app.run()
+    return None
+
+
+if __name__ == "__main__":
+    main()

yunta/dca.py

@@ -0,0 +1,172 @@
+"""Code from Humpreys, Science, 2021 (DOI: 10.1126/science.abm4805), reimplemented for Tensorflow 2.
+
+According to the paper's supplementary material:
+"We reimplemented DCA so that it can be computed using GPUs to speed up the calculation.
+The code is in the box below. In addition, we applied average product correction (APC).[...]"
+
+```python
+import sys
+import numpy as np
+import string
+import tensorflow as tf
+
+def parse_a3m(a3mlines):
+    seqs = []
+    labels = []
+    for line in a3mlines:
+        if line[0] == '>':
+            labels.append(line.rstrip())
+        else:
+            seqs.append(line[:-1])
+    alphabet = np.array(list("ARNDCQEGHILKMFPSTWYV-"), dtype='|S1').view(np.uint8)
+    seq_num = np.array([list(s) for s in seqs], dtype='|S1').view(np.uint8)
+    for i in range(alphabet.shape[0]):
+        seq_num[seq_num == alphabet[i]] = i
+    seq_num[seq_num > 20] = 20
+    return {'seqs' : seq_num, 'labels' : labels }
+
+def tf_cov(x,w=None):
+    if w is None:
+        num_points = tf.cast(tf.shape(x)[0],tf.float32) - 1
+        x_mean = tf.reduce_mean(x, axis=0, keep_dims=True)
+        x = (x - x_mean)
+    else:
+        num_points = tf.reduce_sum(w) - tf.sqrt(tf.reduce_mean(w))
+        x_mean = tf.reduce_sum(x * w[:,None], axis=0, keepdims=True) / num_points
+        x = (x - x_mean) * tf.sqrt(w[:,None])
+    return tf.matmul(tf.transpose(x),x)/num_points
+
+fp = open(sys.argv[1] + ".alignments", "r")
+pair2a3m = {}
+pair = ""
+pairs = []
+a3m = []
+for line in fp:
+    if line[:2] == ">>":
+        if pair and a3m:
+            pair2a3m[pair] = a3m
+            pairs.append(pair)
+        pair = line[2:-1]
+        a3m = []
+    else:
+        a3m.append(line)
+fp.close()
+if pair and a3m:
+    pair2a3m[pair] = a3m
+    pairs.append(pair)
+
+config = tf.ConfigProto(gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.9))
+
+with tf.Graph().as_default():
+    x = tf.placeholder(tf.uint8,shape=(None,None),name="x")
+    x_shape = tf.shape(x)
+    x_nr = x_shape[0]
+    x_nc = x_shape[1]
+    x_ns = 21
+    x_msa = tf.one_hot(x,x_ns)
+    x_cutoff = tf.cast(x_nc,tf.float32) * 0.8
+    x_pw = tf.tensordot(x_msa, x_msa, [[1,2], [1,2]])
+    x_cut = x_pw > x_cutoff
+    x_weights = 1.0/tf.reduce_sum(tf.cast(x_cut, dtype=tf.float32),-1)
+    x_feat = tf.reshape(x_msa,(x_nr,x_nc*x_ns))
+    x_c = tf_cov(x_feat,x_weights) + tf.eye(x_nc*x_ns) *
+    4.5/tf.sqrt(tf.reduce_sum(x_weights))
+    x_c_inv = tf.linalg.inv(x_c)
+    x_w = tf.reshape(x_c_inv,(x_nc,x_ns,x_nc,x_ns))
+    x_wi = tf.sqrt(tf.reduce_sum(tf.square(x_w[:,:-1,:,:-1]),(1,3))) * (1-tf.eye(x_nc))
+    # APC
+    #x_ap = tf.reduce_sum(x_wi,0,keepdims=True) * tf.reduce_sum(x_wi,1,keepdims=True) /
+    tf.reduce_sum(x_wi)
+    #x_wip = (x_wi - x_ap) * (1-tf.eye(x_nc))
+
+    with tf.Session(config=config) as sess:
+        results = []
+        get_pairs = []
+        for pair in pairs:
+            print (pair)
+            msa = parse_a3m(pair2a3m[pair])
+            try:
+                wip = sess.run(x_wi,{x:msa['seqs']})
+                results.append(wip.astype(np.float16))
+                get_pairs.append(pair)
+            except tf.errors.ResourceExhaustedError as e:
+                pass
+        np.savez_compressed(sys.argv[1], *results, names=get_pairs)
+        rp = open(sys.argv[1] + ".log","w")
+        for pair in get_pairs:
+            rp.write(pair + "\n")
+        rp.close()
+```
+"""
+
+from typing import Optional, Union
+from io import TextIOWrapper
+import sys
+
+from carabiner import print_err
+import numpy as np
+import tensorflow as tf
+
+from .structs.msa import MSA, _A3M_ALPHABET
+
+@tf.function
+def _tf_cov(x: tf.Tensor, w: Optional[tf.Tensor] = None) -> tf.float32:
+    if w is None:
+        num_points = tf.cast(tf.shape(x)[0], tf.float32) - 1.
+        x_mean = tf.math.reduce_mean(x, axis=0, keep_dims=True)
+        x = (x - x_mean)
+    else:
+        num_points = tf.math.reduce_sum(w) - tf.math.sqrt(tf.math.reduce_mean(w))
+        x_mean = tf.math.reduce_sum(x * w[:,tf.newaxis], axis=0, keepdims=True) / num_points
+        x = (x - x_mean) * tf.math.sqrt(w[:,tf.newaxis])
+    return tf.linalg.matmul(x, x, transpose_a=True) / num_points
+
+
+@tf.function
+def _get_wip(x: tf.Tensor, apc: bool = False) -> tf.float32:
+    alphabet_size = len(_A3M_ALPHABET)
+    n_row, n_col = tf.shape(x)[-2], tf.shape(x)[-1]
+    x_cutoff = tf.cast(n_col, dtype=tf.float32) * .8
+    msa_one_hot = tf.one_hot(x, alphabet_size)
+    x_pw = tf.tensordot(msa_one_hot, msa_one_hot, [[1,2], [1,2]])
+    x_cut = tf.cast(x_pw > x_cutoff, 
+                    dtype=tf.float32)
+
+    weights = 1. / tf.math.reduce_sum(x_cut, axis=-1)
+    x_feat = tf.reshape(msa_one_hot, (n_row, n_col * alphabet_size))
+
+    x_c = (_tf_cov(x_feat, weights) + tf.eye(n_col * alphabet_size) * 4.5 / 
+           tf.sqrt(tf.math.reduce_sum(weights)))
+    x_c_inv = tf.linalg.inv(x_c)
+
+    x_w = tf.reshape(x_c_inv, (n_col, alphabet_size, n_col, alphabet_size))
+    x_wi = tf.sqrt(tf.reduce_sum(tf.square(x_w[:,:-1,:,:-1]), (1, 3))) * (1. - tf.eye(n_col))
+
+    if apc:
+        x_ap = (tf.math.reduce_sum(x_wi, axis=0, keepdims=True) * tf.math.reduce_sum(x_wi, axis=1, keepdims=True) 
+                / tf.math.reduce_sum(x_wi))
+        x_wi = (x_wi - x_ap) * (1. - tf.eye(n_col))
+
+    return x_wi
+
+
+def calculate_dca(msa: MSA, 
+                  apc: bool = False,
+                  gpu: bool = True) -> np.ndarray:
+
+    """
+    
+    """
+    print_err(f"Devices available:\n{tf.config.list_physical_devices()}")
+    if gpu:
+        gpus = tf.config.list_physical_devices('GPU')
+        if len(gpus) == 0:
+            print_err("WARNING! GPU requested but none available. Falling back to CPU.")
+    msa_token_ids = np.asarray(msa.sequence_token_ids)
+    try:
+        wip = _get_wip(msa_token_ids, apc=apc)
+    except tf.errors.ResourceExhaustedError as e:
+        with tf.device('/cpu:0'):
+            wip = _get_wip(msa_token_ids, apc=apc)
+
+    return wip.numpy()

yunta/dca_torch.py

@@ -0,0 +1,99 @@
+"""DCA for GPU implemented in pytorch."""
+
+from typing import Optional, Union
+
+from io import TextIOWrapper
+import sys
+
+from carabiner import print_err
+import numpy as np
+import torch
+from torch import FloatTensor, Tensor
+import torch.nn.functional as F
+
+from .structs.msa import MSA, _A3M_ALPHABET
+
+
+def _torch_cov(x: Tensor, w: Optional[Tensor] = None) -> FloatTensor:
+    if w is None:
+        return torch.cov(x)
+    else:
+        num_points = torch.sum(w) - torch.sqrt(torch.mean(w))
+        x_mean = torch.sum(x * w.unsqueeze(-1), 
+                           dim=0, keepdim=True) / num_points
+        x = (x - x_mean) * torch.sqrt(w.unsqueeze(-1))
+        return torch.matmul(x.transpose(-2, -1), x) / num_points
+
+
+def _get_wip(x: Tensor, apc: bool = False, gpu: bool = True, dtype=torch.float16) -> FloatTensor:
+    device = torch.device("cuda" if( torch.cuda.is_available() and gpu) else "cpu")
+    # x = x.to(device)
+    alphabet_size = torch.tensor(len(_A3M_ALPHABET), device=device)
+    n_row, n_col = x.shape
+    msa_one_hot = (F.one_hot(x.to(torch.int64), 
+                             num_classes=alphabet_size)
+                    .to(dtype))  # nrow, ncol, alphabet_size
+    identity_counts = torch.tensordot(msa_one_hot, msa_one_hot, 
+                                      dims=[[1,2], [1,2]])  # nrow, nrow
+    # # Should be ~alphabet_size-fold fewer calculations and lower memory
+    # # but actually appears to be slower?
+    # identity_counts = (torch.eq(x, x)
+    #                    .to(dtype)
+    #                    .sum(dim=-1))  # nrow, nrow
+    identity_cutoff = torch.tensor(n_col * .8, device=device)
+    x_cut = (identity_counts > identity_cutoff).to(dtype)
+
+    weights = 1. / torch.sum(x_cut, dim=-1)  # nrow
+    msa_concat_one_hot = msa_one_hot.view(n_row, n_col * alphabet_size)
+
+    shrinkage_coeff = torch.eye(msa_concat_one_hot.shape[-1], 
+                                device=device) * (4.5 / torch.sqrt(torch.sum(weights)))
+    covariance_matrix = _torch_cov(msa_concat_one_hot, weights)  # nrow, nrow
+    shrunk_cov_matrix = covariance_matrix + shrinkage_coeff
+    shrunk_cov_matrix_inv = torch.linalg.inv(shrunk_cov_matrix).view(n_col, alphabet_size, n_col, alphabet_size)
+
+    shrunk_cov_matrix_inv_no_gaps = shrunk_cov_matrix_inv[:,:-1,:,:-1]
+    I_ncol = torch.eye(n_col, device=device)
+    interchain_scores = torch.sqrt(torch.sum(torch.square(shrunk_cov_matrix_inv_no_gaps), 
+                                             dim=(1, 3))) * (1. - I_ncol)
+
+    if apc:
+        apc_factor = (torch.sum(interchain_scores, 
+                                dim=0, keepdim=True)
+                      * torch.sum(interchain_scores, 
+                                  dim=1, keepdim=True) 
+                      / torch.sum(interchain_scores))
+        interchain_scores = (interchain_scores - apc_factor) * (1. - I_ncol)
+
+    return interchain_scores
+
+
+def calculate_dca(msa: MSA, 
+                  apc: bool = False,
+                  gpu: bool = True) -> np.ndarray:
+
+    """
+    
+    """
+    n_gpu = torch.cuda.device_count()
+    print_err(f"GPUs available: {n_gpu}")
+    if gpu:
+        if n_gpu == 0:
+            print_err("WARNING! GPU requested but none available. Falling back to CPU.")
+            device = 'cpu'
+        else:
+            device = 'cuda'
+            # torch.cuda.empty_cache()
+    else:
+        device = 'cpu'
+    msa_token_ids = torch.tensor(msa.sequence_token_ids,
+                                 dtype=torch.int64,
+                                 device=torch.device(device))
+    
+    try:
+        wip = _get_wip(msa_token_ids, apc=apc, gpu=device == 'cuda')
+    except torch.cuda.OutOfMemoryError as e:
+        print_err("GPU memory exhausted; falling back to CPU.")
+        wip = _get_wip(msa_token_ids.to('cpu'), apc=apc, gpu=False)
+
+    return wip.detach().cpu().numpy()

yunta/io.py

@@ -0,0 +1,69 @@
+"""Tools for input and output."""
+
+from typing import Any, Dict, Iterable, Optional, Tuple, Union
+
+from csv import DictWriter
+from dataclasses import is_dataclass, asdict
+from io import TextIOWrapper
+import os
+import shutil
+import tempfile
+
+from bioino import FastaCollection
+from carabiner import print_err
+from carabiner.cast import cast
+import numpy as np
+from pandas import DataFrame
+
+from .structs.pdb_structs import ATMRecord
+from .structs.metrics import ModelMetrics
+
+def save_design(pdb_info,
+                output_name: str, 
+                chainA_length: int) -> None:
+
+    """Save the resulting protein-peptide design to a pdb file.
+
+    """
+
+    chain_name = 'A'
+    with open(output_name, 'w') as f:
+        pdb_contents = pdb_info.split('\n')
+        for line in pdb_contents:
+            try:
+                record = ATMRecord(line)
+                if record.res_no > chainA_length:
+                    chain_name = 'B'
+                outline = f"{line[:21]}{chain_name}{line[22:]}"
+                print(outline, file=f)
+            except Exception:
+                print(line, file=f)
+    return None
+
+
+def write_metrics(metrics: Union[Iterable, Any],
+                  filename: Union[str, TextIOWrapper], 
+                  mode: str = 'w'):
+    
+    """
+    
+    """
+
+    if is_dataclass(metrics):
+        metrics = [metrics]
+    if not all(is_dataclass(m) for m in metrics):
+        raise TypeError("All metrics must be dataclass objects.")
+    f = cast(filename, to=TextIOWrapper, mode=mode)
+    if f.name != '<stdout>':
+        dirname = os.path.dirname(f.name)
+        if len(dirname) > 0  and not os.path.exists(dirname):
+            print_err(f"Creating output directory {dirname}")
+            os.makedirs(dirname)
+
+    w = DictWriter(f, fieldnames=list(asdict(metrics[0])), delimiter='\t')
+    if mode == 'w':
+        w.writeheader()
+    for metric in metrics:
+        w.writerow(asdict(metric))
+
+    return None

yunta/modelling.py

@@ -0,0 +1,43 @@
+"""Tools for setting up and using models."""
+
+from typing import Optional
+
+import os
+from time import time
+
+from carabiner import print_err
+
+from .weights import get_model_weights
+
+_params_path = os.path.join(os.path.dirname(__file__), "data")
+
+def make_model_runner(num_ensemble: int = 1,
+                      max_recycles: int = 10,
+                      param_dir: Optional[str] = None,
+                      model_name: Optional[str] = None):
+    
+    """Generate an AlphaFold2 model runner.
+
+    """
+    
+    import tensorflow as tf
+    tf.config.experimental.set_visible_devices([], "GPU")    
+    from .src_speedppi.alphafold.model import config, data, model
+    from jax.lib import xla_bridge
+
+    print_err(f"Setting up AlphaFold2 model. XLA platform available: {xla_bridge.get_backend().platform}")
+
+    if model_name is None:
+        model_name = 'model_1'
+    if param_dir is None:
+        param_dir = get_model_weights(model_name)
+
+    model_config = config.model_config(model_name)
+    model_config.data.eval.num_ensemble = num_ensemble
+    model_config.data.common.num_recycle = max_recycles
+    model_config.model.num_recycle = max_recycles
+
+    model_params = data.get_model_haiku_params(model_name=model_name, 
+                                               data_dir=param_dir)
+
+    return model.RunModel(model_config, model_params)

yunta/plots.py

@@ -0,0 +1,55 @@
+"""Plotting functions."""
+
+from typing import Optional
+
+import os
+
+from carabiner import colorblind_palette, print_err
+from carabiner.mpl import grid
+from pandas import DataFrame
+import numpy as np
+
+def plot_matrix(m: np.ndarray, 
+                filename_prefix: Optional[str] = None,
+                format: str = 'png',
+                dpi: int = 300, 
+                vline: Optional[float] = None, 
+                hline: Optional[float] = None,
+                vmax: Optional[float] = None,
+                *args, **kwargs):
+    
+    fig, axes = grid()
+    im = axes.imshow(m, cmap='magma', vmin=0., vmax=vmax)
+    fig.colorbar(im, shrink=.7)
+    if hline is not None:
+        axes.axhline(hline, color='lightgrey', zorder=10)
+    if vline is not None:
+        axes.axvline(vline, color='lightgrey', zorder=10)
+    axes.set(*args, **kwargs)
+
+    if filename_prefix is not None:
+        data_file = f"{filename_prefix}.tsv"
+        plot_dir = os.path.dirname(data_file)
+        if not os.path.exists(plot_dir):
+            print_err(f"Creating output directory {plot_dir}")
+            os.makedirs(plot_dir)
+        print_err(f"Saving matrix data as {data_file}")
+        (DataFrame(m, 
+                   index=np.arange(m.shape[0]), 
+                   columns=np.arange(m.shape[1]))
+         .to_csv(data_file, sep='\t', index=False))
+        plot_file = f"{filename_prefix}.{format}"
+        print_err(f"Saving matrix plot as {plot_file}")
+        fig.savefig(plot_file, bbox_inches="tight", dpi=dpi)
+
+    return fig, axes
+
+def plot_dca(dca: np.ndarray, 
+             filename_prefix: Optional[str] = None,
+             format: str = 'png',
+             dpi: int = 300):
+    
+    if filename_prefix is not None:
+        filename_prefix += "_dca"
+
+    return plot_matrix(dca, filename_prefix=filename_prefix, format=format, dpi=dpi)