Coverage for rest_api/data_entry/sample_import.py: 32%
139 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-10 03:22 +0000
1from dataclasses import dataclass
2import pathlib
3import pickle
4import typing
5from typing import Any
6from typing import Optional
8from django.utils import timezone
9import pandas as pd
11from rest_api.models import Alignment
12from rest_api.models import AminoAcidMutation
13from rest_api.models import CDS
14from rest_api.models import Gene
15from rest_api.models import NucleotideMutation
16from rest_api.models import Replicon
17from rest_api.models import Sample
18from rest_api.models import Sequence
19from sonar_backend.settings import LOGGER
22@dataclass
23class SampleRaw:
24 anno_vcf_file: str
25 cds_file: str
26 header: str
27 name: str
28 properties: dict
29 ref_file: str
30 refmol: str
31 refmolid: int
32 seq_file: str
33 seqhash: str
34 sequence_length: int
35 sourceid: int
36 translationid: int
37 include_nx: bool
38 var_parquet_file: Optional[str] = None
39 vcffile: Optional[str] = None
40 algnid: Optional[int] = None
41 sequenceid: Optional[int] = None
42 refseq_id: Optional[int] = None
43 source_acc: Optional[str] = None
44 algn_file: Optional[str] = None
45 anno_tsv_file: Optional[str] = None
46 lift_file: Optional[str] = None
49@dataclass
50class VarRaw:
51 id: int
52 ref: str
53 start: int
54 end: int
55 alt: str | None
56 replicon_or_cds_accession: str
57 # lable: str
58 type: str
59 frameshift: int
60 parent_id: list[int] | None
63class VCFInfoLOFRaw:
64 # Predicted loss of function effects for this variant.
65 gene_name: str
66 gene_id: str
67 number_of_transcripts_in_gene: str
68 percent_of_transcripts_affected: str
71class VCFInfoNMDRaw:
72 # Predicted nonsense mediated decay effects for this variant.
73 gene_name: str
74 gene_id: str
75 number_of_transcripts_in_gene: str
76 percent_of_transcripts_affected: str
79class SonarImport:
80 def __init__(
81 self,
82 path: pathlib.Path,
83 import_folder="import_data",
84 ):
85 self.sample_file_path = path
86 self.import_folder = import_folder
87 self.sample_raw = SampleRaw(**self._import_pickle(path))
88 self.sequence: None | Sequence = None
89 self.sample: None | Sample = None
90 self.replicon: None | Replicon = None
91 self.alignment: None | Alignment = None
92 self.success = False
94 if self.sample_raw.var_parquet_file:
95 self.vars_raw = [
96 var
97 for var in self._import_vars(
98 self.sample_raw.var_parquet_file, self.sample_raw.include_nx
99 )
100 ]
101 else:
102 raise Exception("No var file found")
104 def get_sample_name(self):
105 return self.sample_raw.name
107 def get_sequence_obj(self):
108 """
109 Create a Sequence object WITHOUT querying or persisting to DB.
110 Bulk upsert will handle create/update efficiently in a single query.
111 """
112 # Always create new object - bulk_create with update_conflicts will handle upsert
113 self.sequence = Sequence(
114 name=self.sample_raw.name,
115 seqhash=self.sample_raw.seqhash,
116 length=self.sample_raw.sequence_length,
117 last_update_date=timezone.now(),
118 )
119 return self.sequence
121 def update_replicon_obj(self, replicon_cache: dict[str, Replicon]):
122 if self.sample_raw.source_acc not in replicon_cache:
123 replicon_cache[self.sample_raw.source_acc] = Replicon.objects.get(
124 accession=self.sample_raw.source_acc
125 )
126 self.replicon = replicon_cache[self.sample_raw.source_acc]
128 def create_alignment(self, alignments_list: list[Alignment]):
129 self.alignment = next(
130 filter(
131 lambda x: x.sequence == self.sequence and x.replicon == self.replicon,
132 alignments_list,
133 ),
134 None,
135 )
136 if not self.alignment:
137 self.alignment = Alignment(sequence=self.sequence, replicon=self.replicon)
138 alignments_list.append(self.alignment)
140 def get_mutation_objs_nt(
141 self,
142 nt_mutation_set: list[NucleotideMutation],
143 replicon_cache: dict[str, Replicon | None],
144 gene_cache_by_var_pos: dict[Replicon | None, dict[int, dict[int, Gene | None]]],
145 nt_mutation_alignment_relations: list[NucleotideMutation.alignments.through],
146 ) -> dict[int, NucleotideMutation]:
147 import_id_to_sample_mutations: dict[int, NucleotideMutation] = {}
148 for var_raw in self.vars_raw:
149 replicon = None
150 if var_raw.type == "nt":
151 if not var_raw.replicon_or_cds_accession in replicon_cache:
152 replicon_cache[var_raw.replicon_or_cds_accession] = (
153 Replicon.objects.get(
154 accession=var_raw.replicon_or_cds_accession
155 )
156 )
157 replicon = replicon_cache[var_raw.replicon_or_cds_accession]
158 if not replicon in gene_cache_by_var_pos:
159 gene_cache_by_var_pos[replicon] = {}
160 if not var_raw.start in gene_cache_by_var_pos[replicon]:
161 gene_cache_by_var_pos[replicon][var_raw.start] = {}
162 if not var_raw.end in gene_cache_by_var_pos[replicon][var_raw.start]:
163 gene_cache_by_var_pos[replicon][var_raw.start][var_raw.end] = (
164 Gene.objects.filter(
165 replicon=replicon,
166 start__gte=var_raw.start,
167 end__lte=var_raw.end,
168 ).first()
169 )
170 gene = gene_cache_by_var_pos[replicon][var_raw.start][var_raw.end]
171 # in DEL, we dont keep REF in the database.
172 if var_raw.alt is None:
173 var_raw.ref = ""
175 mutation_data = {
176 "ref": var_raw.ref if var_raw.ref else "",
177 "alt": var_raw.alt if var_raw.alt else "",
178 "start": var_raw.start,
179 "end": var_raw.end,
180 "replicon": self.replicon,
181 "is_frameshift": var_raw.frameshift,
182 }
183 mutation = next(
184 filter(
185 lambda x: self.is_same_mutation(mutation_data, x),
186 nt_mutation_set,
187 ),
188 None,
189 )
190 if not mutation:
191 mutation = NucleotideMutation(**mutation_data)
192 nt_mutation_set.append(mutation)
193 nt_mutation_alignment_relations.append(
194 NucleotideMutation.alignments.through(
195 nucleotidemutation=mutation, alignment=self.alignment
196 )
197 )
198 import_id_to_sample_mutations[var_raw.id] = mutation
199 return import_id_to_sample_mutations
201 def is_same_mutation(
202 self, mutation_data: dict, mutation: NucleotideMutation | AminoAcidMutation
203 ) -> bool:
204 return all(getattr(mutation, k) == v for k, v in mutation_data.items())
206 def get_mutation_objs_cds_and_parent_relations(
207 self,
208 cds_mutation_set: list[AminoAcidMutation],
209 gene_cache_by_accession: dict[str, CDS | None],
210 parent_id_mapping: dict[int, NucleotideMutation],
211 aa_mutation_alignment_relations: list[AminoAcidMutation.alignments.through],
212 ) -> list[AminoAcidMutation.parent.through]:
213 sample_cds_mutations: list[AminoAcidMutation] = []
214 mutation_parent_relations = []
215 for var_raw in self.vars_raw:
216 if var_raw.type == "cds":
217 if not var_raw.replicon_or_cds_accession in gene_cache_by_accession:
218 try:
219 gene_cache_by_accession[var_raw.replicon_or_cds_accession] = (
220 CDS.objects.get(accession=var_raw.replicon_or_cds_accession)
221 )
222 except CDS.DoesNotExist:
223 LOGGER.error(
224 f"CDS not found for accession: {var_raw.replicon_or_cds_accession}"
225 )
226 continue
227 cds = gene_cache_by_accession[var_raw.replicon_or_cds_accession]
228 if var_raw.alt is None:
229 var_raw.ref = ""
230 mutation_data = {
231 "cds": cds,
232 "ref": var_raw.ref if var_raw.ref else "",
233 "alt": var_raw.alt if var_raw.alt else "",
234 "start": var_raw.start if var_raw.start else 0,
235 "end": var_raw.end if var_raw.end else 0,
236 }
237 mutation = next(
238 filter(
239 lambda x: self.is_same_mutation(mutation_data, x),
240 cds_mutation_set,
241 ),
242 None,
243 )
244 if not mutation:
245 mutation = AminoAcidMutation(**mutation_data)
246 cds_mutation_set.append(mutation)
247 aa_mutation_alignment_relations.append(
248 AminoAcidMutation.alignments.through(
249 aminoacidmutation=mutation, alignment=self.alignment
250 )
251 )
252 if var_raw.parent_id:
253 # This KeyError occurs due to the skip-nx flag.
254 # If the third nucleotide of a codon is 'N', the corresponding nucleotide mutation is not read in.
255 # However, if the first or second nucleotide is also mutated, an amino acid mutation can still be inferred
256 # due to the ambiguity of the genetic code.
257 # For this mutations we cannot store the relation to the nt parent
258 for parent_id in var_raw.parent_id:
259 try:
260 mutation_parent_relations.append(
261 AminoAcidMutation.parent.through(
262 aminoacidmutation=mutation,
263 nucleotidemutation=parent_id_mapping[parent_id],
264 )
265 )
266 except KeyError:
267 LOGGER.warning(
268 f"Parent ID {parent_id} not found in parent_id_mapping for mutation {mutation_data}"
269 )
270 pass
271 sample_cds_mutations.append(mutation)
272 return mutation_parent_relations
274 def _import_pickle(self, path: str):
275 with open(path, "rb") as f:
276 return pickle.load(f)
278 def _import_vars(self, path, include_nx: bool):
279 file_name = pathlib.Path(path).name
280 self.var_file_path = (
281 pathlib.Path(self.import_folder)
282 .joinpath("var")
283 .joinpath(file_name[:2])
284 .joinpath(file_name)
285 )
286 var_df = pd.read_parquet(self.var_file_path)
287 var_df[["ref", "alt"]] = var_df[["ref", "alt"]].fillna("").replace({" ": ""})
289 if not include_nx:
290 # remove all ref containing Ns for nt, or X for cds
291 var_df = var_df[
292 ~((var_df["type"] == "nt") & var_df["alt"].str.contains("N", na=False))
293 ]
294 var_df = var_df[
295 ~((var_df["type"] == "cds") & var_df["alt"].str.contains("X", na=False))
296 ]
297 for _, row in var_df.iterrows():
298 try:
299 yield VarRaw(
300 row["id"],
301 row["ref"],
302 row["start"],
303 row["end"],
304 row["alt"],
305 row["reference_acc"],
306 row["type"],
307 row["frameshift"],
308 (
309 [int(x) for x in row["parent_id"].split(",")]
310 if pd.notna(row["parent_id"]) and row["parent_id"].strip() != ""
311 else None
312 ),
313 )
314 except Exception as e:
315 print(f"Error processing row: {row}")
316 print(f"Error file: {self.var_file_path}")
317 raise e
319 def _import_seq(self, path):
320 file_name = pathlib.Path(path).name
321 self.seq_file_path = (
322 pathlib.Path(self.import_folder)
323 .joinpath("seq")
324 .joinpath(file_name[:2])
325 .joinpath(file_name)
326 )
327 with open(self.seq_file_path, "r") as handle:
328 for line in handle:
329 yield line.strip("\r\n")