def parse_config(config_str: str, preserve_order: bool = True) -> dict:
"""
Parse a configuration string
Parameters
----------
config_str : str
Configuration to be parsed
preserve_order : bool, optional (default: False)
Preserve formatting of input configuration
string
Returns
-------
dict
Configuration dictionary
"""
try:
if preserve_order:
return yaml.load(config_str, Loader=yaml.RoundTripLoader)
else:
return yaml.safe_load(config_str)
except ParserError as e:
raise InvalidParameterError(
"Could not parse input configuration. "
"Formatting mistake in config file? "
"See ParserError above for details.") from e
def run(**kwargs):
"""
Run embedding protocol
Parameters
----------
kwargs arguments (* denotes optional):
prefix: Output prefix for all generated files
protocol: Which protocol to use
stage_name: The stage name
Returns
-------
Dictionary with results of stage
"""
check_required(kwargs, ['protocol', 'prefix', 'stage_name'])
if kwargs["protocol"] not in PROTOCOLS:
raise InvalidParameterError(
"Invalid protocol selection: " +
"{}. Valid protocols are: {}".format(
kwargs["protocol"], ", ".join(PROTOCOLS.keys())
)
)
return PROTOCOLS[kwargs["protocol"]](**kwargs)
def run(**kwargs):
"""
Align query sequences with target sequences
Parameters
----------
kwargs arguments (* denotes optional):
prefix: Output prefix for all generated files
stage_name: The stage name
protocol: Which method to use for alignment
mapping_file: The mapping file generated by the pipeline when remapping indexes
remapped_sequences_file: The fasta file with entries corresponding to the mapping file
Returns
-------
Dictionary with results of stage
"""
check_required(
kwargs,
[
"protocol",
"prefix",
"stage_name",
"remapped_sequences_file",
"mapping_file",
],
)
if kwargs["protocol"] not in PROTOCOLS:
raise InvalidParameterError(
"Invalid protocol selection: " +
"{}. Valid protocols are: {}".format(kwargs["protocol"], ", ".join(
PROTOCOLS.keys())))
return PROTOCOLS[kwargs["protocol"]](**kwargs)
def read_config_file(config_path: Union[str, Path],
preserve_order: bool = True) -> dict:
"""
Read config from path to file.
:param config_path: path to .yml config file
:param preserve_order:
:return:
"""
with open(config_path, "r") as fp:
try:
if preserve_order:
return yaml.load(fp, Loader=yaml.RoundTripLoader)
else:
return yaml.safe_load(fp)
except YAMLError as e:
raise InvalidParameterError(
f"Could not parse configuration file at '{config_path}' as yaml. "
"Formatting mistake in config file? "
"See Error above for details.") from e
def deepblast(**kwargs) -> Dict[str, Any]:
"""Sequence-Sequence alignments with DeepBLAST
DeepBLAST learned structural alignments from sequence
https://github.com/flatironinstitute/deepblast
https://www.biorxiv.org/content/10.1101/2020.11.03.365932v1
"""
# TODO: Fix that logic before merging
if "transferred_annotations_file" not in kwargs and "pairings_file" not in kwargs:
raise MissingParameterError(
"You need to specify either 'transferred_annotations_file' or 'pairings_file' for DeepBLAST"
)
if "transferred_annotations_file" in kwargs and "pairings_file" in kwargs:
raise InvalidParameterError(
"You can't specify both 'transferred_annotations_file' and 'pairings_file' for DeepBLAST"
)
result_kwargs = deepcopy(kwargs)
file_manager = get_file_manager(**kwargs)
# This stays below 8GB, so it should be a good default
batch_size = result_kwargs.setdefault("batch_size", 50)
if "device" in result_kwargs:
device = torch.device(result_kwargs["device"])
if device.type != "cuda":
raise RuntimeError(
f"You can only run DeepBLAST on a CUDA-compatible GPU, not on {device.type}"
)
else:
if not torch.cuda.is_available():
raise RuntimeError(
"DeepBLAST requires a CUDA-compatible GPU, but none was found")
device = torch.device("cuda")
mapping_file = read_mapping_file(result_kwargs["mapping_file"])
mapping = {
str(remapped): original
for remapped, original in mapping_file[["original_id"]].itertuples()
}
query_by_id = {
mapping[entry.name]: str(entry.seq)
for entry in SeqIO.parse(result_kwargs["remapped_sequences_file"],
"fasta")
}
# You can either provide a set of pairing or use the output of k-nn with a fasta file for the reference embeddings
if "pairings_file" in result_kwargs:
pairings_file = read_csv(result_kwargs["pairings_file"])
pairings = list(pairings_file[["query",
"target"]].itertuples(index=False))
target_by_id = query_by_id
else:
transferred_annotations_file = read_csv(
result_kwargs["transferred_annotations_file"])
pairings = []
for _, row in transferred_annotations_file.iterrows():
query = row["original_id"]
for target in row.filter(regex="k_nn_.*_identifier"):
pairings.append((query, target))
target_by_id = {}
for entry in read_fasta(result_kwargs["reference_fasta_file"]):
target_by_id[entry.name] = str(entry.seq[:])
# Create one output file per query
result_kwargs["alignment_files"] = dict()
for query in set(i for i, _ in pairings):
filename = file_manager.create_file(
result_kwargs.get("prefix"),
result_kwargs.get("stage_name"),
f"{slugify(query, lowercase=False)}_alignments",
extension=".a2m",
)
result_kwargs["alignment_files"][query] = filename
unknown_queries = set(list(zip(*pairings))[0]) - set(query_by_id.keys())
if unknown_queries:
raise ValueError(f"Unknown query sequences: {unknown_queries}")
unknown_targets = set(list(zip(*pairings))[1]) - set(target_by_id.keys())
if unknown_targets:
raise ValueError(f"Unknown target sequences: {unknown_targets}")
# Load the pretrained model
if "model_file" not in result_kwargs:
model_file = get_model_file("deepblast", "model_file")
else:
model_file = result_kwargs["model_file"]
alignments = deepblast_align(pairings, query_by_id, target_by_id,
model_file, device, batch_size)
for query, alignments in itertools.groupby(alignments, key=lambda i: i[0]):
_, targets, queries_aligned, targets_aligned = list(zip(*alignments))
padded_query, padded_targets = pairwise_alignments_to_msa(
queries_aligned, targets_aligned)
with open(result_kwargs["alignment_files"][query], "w") as fp:
fp.write(f">{query}\n")
fp.write(f"{padded_query}\n")
for target, padded_target in zip(targets, padded_targets):
fp.write(f">{target}\n")
fp.write(f"{padded_target}\n")
return result_kwargs