File size: 5,365 Bytes
80a72c3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | """
Class that represents a file with the prediction results
* allows prediction results to be appended to the file in chunks
* allows results to be skipped if they already exist in the file
"""
import csv
from typing import List, OrderedDict
from pathlib import Path
import logging
from .models.results import PredictionResult
from .errors import PredictionResultExistsError, FileExistsError
LOG = logging.getLogger(__name__)
class PredictionResultsFileIter:
def __init__(self, results_file_class):
self._results = results_file_class.get_results()
self._results_size = len(self._results)
self._current_index = 0
def __iter__(self):
return self
def __next__(self):
if self._current_index < self._results_size:
member = self._results[self._current_index]
self._current_index += 1
return member
raise StopIteration
class PredictionResultsFile:
"""
Writes prediction results to a file in chunks
"""
COLNAMES = ['chain_id', 'sequence_md5', 'nres', 'ndom', 'chopping', 'confidence', 'time_sec']
def __init__(self,
csv_path: Path, *,
chunk_size: int=20,
write_header: bool=True,
allow_append: bool=False,
allow_skip: bool=False):
self.csv_path = Path(str(csv_path)).absolute()
self.chunk_size = chunk_size
self.write_header = write_header
self.allow_append = allow_append
self.allow_skip = allow_skip
self._results_by_id: OrderedDict[PredictionResult] = OrderedDict()
self._unflushed_results_by_id: OrderedDict[PredictionResult] = OrderedDict()
self._written_header = False
self._init()
def _init(self):
if self.csv_path.exists():
msg = f"file '{self.csv_path}' already exists (allow_append={self.allow_append}))"
LOG.warning(msg)
if not self.allow_append:
raise FileExistsError(msg)
self._read_results()
def _read_results(self):
"""
Read the results from the file
"""
with self.csv_path.open('r') as fp:
reader = self.get_csv_reader(fp)
for row in reader:
if row['chain_id'] == 'chain_id':
continue
pdb_path = Path(row['chain_id'] + '.pdb')
result = PredictionResult(pdb_path=pdb_path, **row)
self._results_by_id[result.chain_id] = result
def get_flushed_results(self):
self.flush()
return self._results_by_id.values()
def get_csv_reader(self, fp):
return csv.DictReader(fp, delimiter='\t', fieldnames=self.COLNAMES)
def get_csv_writer(self, fp):
return csv.DictWriter(fp, delimiter='\t', fieldnames=self.COLNAMES)
def write_csv_result(self, csv_writer, res: PredictionResult):
"""
Render PredictionResult result to file pointer
"""
row = {
'chain_id': res.chain_id,
'sequence_md5': res.sequence_md5,
'nres': res.nres,
'ndom': res.ndom,
'chopping': res.chopping if res.chopping is not None else 'NULL',
'confidence': f'{res.confidence:.3g}' if res.confidence is not None else 'NULL',
'time_sec': f'{res.time_sec}' if res.time_sec is not None else 'NULL',
}
csv_writer.writerow(row)
def flush(self):
"""
Write unflushed results to the output file
"""
if len(self._unflushed_results_by_id) == 0:
return
with self.csv_path.open('a') as fp:
csv_writer = self.get_csv_writer(fp)
if self.write_header and not self._written_header and len(self._results_by_id) == 0:
csv_writer.writeheader()
self._written_header = True
for result in self._unflushed_results_by_id.values():
self.write_csv_result(csv_writer, result)
def has_result(self, result: PredictionResult):
"""
Check if the prediction result already exists in the file
"""
return self.has_result_for_chain_id(result.chain_id)
def has_result_for_chain_id(self, chain_id: str):
"""
Check if the chain_id already exists in the file
"""
if chain_id in self._results_by_id:
return True
if chain_id in self._unflushed_results_by_id:
return True
return False
def add_result(self, result: PredictionResult):
"""
Add a result to the buffer (and possibly flush the buffer)
"""
if self.has_result(result):
msg = f"result '{result.chain_id}' already exists in file '{self.csv_path}'"
raise PredictionResultExistsError(msg)
self._unflushed_results_by_id[result.chain_id] = result
if len(self._unflushed_results_by_id) >= self.chunk_size:
self.flush()
self._results_by_id.update(self._unflushed_results_by_id)
self._unflushed_results_by_id = OrderedDict()
def add_results(self, results: List[PredictionResult]):
for res in results:
self.add_result(res)
def __iter__(self):
return PredictionResultsFileIter(self)
|