From e442d97ec3556ffd141eeef502bcf3d620f8d328 Mon Sep 17 00:00:00 2001 From: kathryn1995 Date: Fri, 2 Oct 2026 21:53:35 -0700 Subject: [PATCH] Add PMOUpdater methods for nested fields from long-format tables --- src/pmotools/pmo_builder/pmo_updater.py | 446 +++++++++++++++++++-- tests/test_pmo_builder/test_pmo_updater.py | 307 +++++++++++++- 2 files changed, 703 insertions(+), 50 deletions(-) diff --git a/src/pmotools/pmo_builder/pmo_updater.py b/src/pmotools/pmo_builder/pmo_updater.py index c0f9305..c47c474 100644 --- a/src/pmotools/pmo_builder/pmo_updater.py +++ b/src/pmotools/pmo_builder/pmo_updater.py @@ -60,28 +60,24 @@ def update_specimen_meta_with_traveler_info( :param replace_current_traveler_info: whether to replace current travel info :return: a reference to the updated PMO """ - required_cols = [ - specimen_name_col, - travel_country_col, - travel_start_col, - travel_end_col, - ] - if bed_net_usage_col is not None: - required_cols.append(bed_net_usage_col) - if geo_admin1_col is not None: - required_cols.append(geo_admin1_col) - if geo_admin2_col is not None: - required_cols.append(geo_admin2_col) - if geo_admin3_col is not None: - required_cols.append(geo_admin3_col) - if lat_lon_col is not None: - required_cols.append(lat_lon_col) - - if not set(required_cols).issubset(traveler_info.columns): - raise Exception( - "missing traveler_info columns: " + ",".join(required_cols), - " columns in table: " + ",".join(traveler_info.columns), - ) + optional_cols = { + "geo_admin1": geo_admin1_col, + "geo_admin2": geo_admin2_col, + "geo_admin3": geo_admin3_col, + "lat_lon": lat_lon_col, + } + PMOUpdater._check_columns( + traveler_info, + [ + specimen_name_col, + travel_country_col, + travel_start_col, + travel_end_col, + bed_net_usage_col, + *optional_cols.values(), + ], + "traveler_info", + ) specimen_names_in_pmo = set(PMOProcessor.get_specimen_names(pmo)) specimen_names_in_traveler_info = set( @@ -95,7 +91,6 @@ def update_specimen_meta_with_traveler_info( raise ValueError( f"Provided traveler info for the following specimens but they are missing from the PMO: {sorted(missing_traveler_specs)}" ) - traveler_info_records = traveler_info[required_cols].to_dict(orient="records") spec_indexs = PMOProcessor.get_index_key_of_specimen_names(pmo) # prep traveler info lists, clear the list if we are replacing or start an empty list to append to if none exist already @@ -109,12 +104,20 @@ def update_specimen_meta_with_traveler_info( "travel_out_six_month" ] = [] - for travel_rec in traveler_info_records: - specimen_name = str(travel_rec[specimen_name_col]) + for _, row in traveler_info.iterrows(): + specimen_name = str(row[specimen_name_col]) + if PMOUpdater._is_blank(row[travel_country_col]): + raise ValueError( + f"Missing required value in column '{travel_country_col}' for specimen '{specimen_name}'" + ) + travel_rec = {"travel_country": str(row[travel_country_col])} # Validate date formats - for date_col in (travel_start_col, travel_end_col): - val = travel_rec[date_col] - if pd.isna(val): + for date_col, key in ( + (travel_start_col, "travel_start_date"), + (travel_end_col, "travel_end_date"), + ): + val = row[date_col] + if PMOUpdater._is_blank(val): raise ValueError( f"Missing required date value in column '{date_col}' for specimen '{specimen_name}'" ) @@ -124,13 +127,396 @@ def update_specimen_meta_with_traveler_info( f"Invalid date format in '{date_col}' for specimen '{specimen_name}': '{val_str}'. " f"Expected YYYY-MM or YYYY-MM-DD" ) - # add in travel_rec - travel_rec.pop(specimen_name_col, None) + travel_rec[key] = val_str + if bed_net_usage_col and not PMOUpdater._is_blank(row[bed_net_usage_col]): + travel_rec["bed_net_usage"] = float(row[bed_net_usage_col]) + for key, col in optional_cols.items(): + if col and not PMOUpdater._is_blank(row[col]): + travel_rec[key] = str(row[col]) pmo["specimen_info"][spec_indexs[specimen_name]][ "travel_out_six_month" ].append(travel_rec) return pmo + @staticmethod + def _is_blank(value) -> bool: + if value is None: + return True + if isinstance(value, str): + return not value.strip() + return bool(pd.isna(value)) + + @staticmethod + def _check_columns(table: pd.DataFrame, cols: list, table_label: str) -> None: + missing = [col for col in cols if col is not None and col not in table.columns] + if missing: + raise ValueError( + f"Missing {table_label} columns: {missing}. Columns in table: {list(table.columns)}" + ) + + @staticmethod + def _genomic_location_from_row( + row, + context: str, + genome_id_col: str | None, + chrom_col: str, + start_col: str, + end_col: str, + strand_col: str | None = None, + ref_seq_col: str | None = None, + alt_seq_col: str | None = None, + ) -> dict: + """ + Build a GenomicLocation dict from a table row. If genome_id_col is None the genome_id defaults to 0. + """ + required = [chrom_col, start_col, end_col] + if genome_id_col is not None: + required.append(genome_id_col) + for col in required: + if PMOUpdater._is_blank(row[col]): + raise ValueError( + f"Missing required value in column '{col}' for {context}" + ) + try: + location = { + "genome_id": int(row[genome_id_col]) if genome_id_col else 0, + "chrom": str(row[chrom_col]), + "start": int(row[start_col]), + "end": int(row[end_col]), + } + except (TypeError, ValueError) as e: + raise ValueError( + f"genome id, start and end must be integers for {context}" + ) from e + for key, col in ( + ("strand", strand_col), + ("ref_seq", ref_seq_col), + ("alt_seq", alt_seq_col), + ): + if col and not PMOUpdater._is_blank(row[col]): + location[key] = str(row[col]) + return location + + @staticmethod + def _get_representative_microhaplotype( + pmo, target_name: str, seq: str, target_indexes: dict + ) -> dict: + if target_name not in target_indexes: + raise ValueError( + f"Target '{target_name}' has no representative microhaplotypes in the PMO" + ) + microhaplotypes = pmo["representative_microhaplotypes"]["targets"][ + target_indexes[target_name] + ]["microhaplotypes"] + for microhaplotype in microhaplotypes: + if microhaplotype["seq"] == seq: + return microhaplotype + raise ValueError( + f"No representative microhaplotype with seq '{seq}' for target '{target_name}'" + ) + + @staticmethod + def update_target_info_with_markers_of_interest( + pmo, + markers_info: pd.DataFrame, + target_name_col: str = "target_name", + chrom_col: str = "chrom", + start_col: str = "start", + end_col: str = "end", + genome_id_col: str | None = None, + strand_col: str | None = None, + ref_seq_col: str | None = None, + alt_seq_col: str | None = None, + associations_col: str | None = None, + associations_delim: str = ",", + replace_current_markers: bool = False, + ): + """ + Update a PMO's target_info with markers of interest, from a table with one row per marker + + :param pmo: the PMO to update, will directly modify this PMO + :param markers_info: the markers of interest table + :param target_name_col: the column containing the name of the target the marker is covered by + :param chrom_col: the column containing the chromosome of the marker + :param start_col: the column containing the start of the marker, 0-based + :param end_col: the column containing the end of the marker, 0-based + :param genome_id_col: (Optional) the column containing the index into targeted_genomes. Default: 0 for every marker + :param strand_col: (Optional) the column containing the strand of the marker + :param ref_seq_col: (Optional) the column containing the reference sequence of the marker + :param alt_seq_col: (Optional) the column containing an alternative sequence of the marker + :param associations_col: (Optional) the column containing associations with the marker, e.g. SP resistance + :param associations_delim: the delimiter between associations. Default: ',' + :param replace_current_markers: whether to replace current markers of interest for the targets in the table + :return: a reference to the updated PMO + """ + PMOUpdater._check_columns( + markers_info, + [ + target_name_col, + chrom_col, + start_col, + end_col, + genome_id_col, + strand_col, + ref_seq_col, + alt_seq_col, + associations_col, + ], + "markers_info", + ) + target_indexes = PMOProcessor.get_index_key_of_target_names(pmo) + target_names_in_table = set(markers_info[target_name_col].astype(str)) + missing_targets = target_names_in_table - set(target_indexes) + if missing_targets: + raise ValueError( + f"Provided markers of interest for the following targets but they are missing from the PMO: {sorted(missing_targets)}" + ) + + for target_name in target_names_in_table: + target = pmo["target_info"][target_indexes[target_name]] + if replace_current_markers or not target.get("markers_of_interest"): + target["markers_of_interest"] = [] + + for _, row in markers_info.iterrows(): + target_name = str(row[target_name_col]) + marker = { + "marker_location": PMOUpdater._genomic_location_from_row( + row, + f"marker of interest for target '{target_name}'", + genome_id_col, + chrom_col, + start_col, + end_col, + strand_col, + ref_seq_col, + alt_seq_col, + ) + } + if associations_col and not PMOUpdater._is_blank(row[associations_col]): + marker["associations"] = [ + association.strip() + for association in str(row[associations_col]).split( + associations_delim + ) + if association.strip() + ] + pmo["target_info"][target_indexes[target_name]][ + "markers_of_interest" + ].append(marker) + return pmo + + @staticmethod + def update_representative_microhaplotypes_with_seq_variants( + pmo, + seq_variants_info: pd.DataFrame, + target_name_col: str = "target_name", + seq_col: str = "seq", + chrom_col: str = "chrom", + start_col: str = "start", + end_col: str = "end", + genome_id_col: str | None = None, + strand_col: str | None = None, + ref_seq_col: str | None = None, + alt_seq_col: str | None = None, + replace_current_seq_variants: bool = False, + ): + """ + Update a PMO's representative microhaplotypes with associated sequence variants, from a table with one row per variant + + :param pmo: the PMO to update, will directly modify this PMO + :param seq_variants_info: the sequence variants table + :param target_name_col: the column containing the target name of the microhaplotype + :param seq_col: the column containing the sequence of the microhaplotype + :param chrom_col: the column containing the chromosome of the variant + :param start_col: the column containing the start of the variant, 0-based + :param end_col: the column containing the end of the variant, 0-based + :param genome_id_col: (Optional) the column containing the index into targeted_genomes. Default: 0 for every variant + :param strand_col: (Optional) the column containing the strand of the variant + :param ref_seq_col: (Optional) the column containing the reference sequence of the variant + :param alt_seq_col: (Optional) the column containing the alternative sequence of the variant + :param replace_current_seq_variants: whether to replace current sequence variants for the microhaplotypes in the table + :return: a reference to the updated PMO + """ + PMOUpdater._check_columns( + seq_variants_info, + [ + target_name_col, + seq_col, + chrom_col, + start_col, + end_col, + genome_id_col, + strand_col, + ref_seq_col, + alt_seq_col, + ], + "seq_variants_info", + ) + target_indexes = ( + PMOProcessor.get_index_key_of_target_in_representative_microhaplotypes(pmo) + ) + microhaplotypes = [ + PMOUpdater._get_representative_microhaplotype( + pmo, str(row[target_name_col]), str(row[seq_col]), target_indexes + ) + for _, row in seq_variants_info.iterrows() + ] + + for microhaplotype in microhaplotypes: + if replace_current_seq_variants or not microhaplotype.get( + "associated_seq_variants" + ): + microhaplotype["associated_seq_variants"] = [] + + for (_, row), microhaplotype in zip( + seq_variants_info.iterrows(), microhaplotypes + ): + microhaplotype["associated_seq_variants"].append( + PMOUpdater._genomic_location_from_row( + row, + f"sequence variant for target '{row[target_name_col]}'", + genome_id_col, + chrom_col, + start_col, + end_col, + strand_col, + ref_seq_col, + alt_seq_col, + ) + ) + return pmo + + @staticmethod + def update_representative_microhaplotypes_with_protein_variants( + pmo, + protein_variants_info: pd.DataFrame, + target_name_col: str = "target_name", + seq_col: str = "seq", + transcript_col: str = "transcript", + protein_start_col: str = "protein_start", + protein_end_col: str = "protein_end", + protein_genome_id_col: str | None = None, + protein_ref_seq_col: str | None = None, + protein_alt_seq_col: str | None = None, + gene_name_col: str | None = None, + alternative_gene_name_col: str | None = None, + codon_chrom_col: str | None = None, + codon_start_col: str | None = None, + codon_end_col: str | None = None, + codon_genome_id_col: str | None = None, + codon_strand_col: str | None = None, + codon_ref_seq_col: str | None = None, + codon_alt_seq_col: str | None = None, + replace_current_protein_variants: bool = False, + ): + """ + Update a PMO's representative microhaplotypes with associated protein variants, from a table with one row per variant + + :param pmo: the PMO to update, will directly modify this PMO + :param protein_variants_info: the protein variants table + :param target_name_col: the column containing the target name of the microhaplotype + :param seq_col: the column containing the sequence of the microhaplotype + :param transcript_col: the column containing the transcript name, used as the chrom of the protein location + :param protein_start_col: the column containing the start of the variant within the protein, 0-based + :param protein_end_col: the column containing the end of the variant within the protein, 0-based + :param protein_genome_id_col: (Optional) the column containing the index into targeted_genomes for the protein location. Default: 0 + :param protein_ref_seq_col: (Optional) the column containing the reference amino acid(s) + :param protein_alt_seq_col: (Optional) the column containing the alternative amino acid(s) + :param gene_name_col: (Optional) the column containing the gene name + :param alternative_gene_name_col: (Optional) the column containing an alternative gene name + :param codon_chrom_col: (Optional) the column containing the chromosome of the codon. Must be set with codon_start_col and codon_end_col + :param codon_start_col: (Optional) the column containing the genomic start of the codon, 0-based + :param codon_end_col: (Optional) the column containing the genomic end of the codon, 0-based + :param codon_genome_id_col: (Optional) the column containing the index into targeted_genomes for the codon location. Default: 0 + :param codon_strand_col: (Optional) the column containing the strand of the codon + :param codon_ref_seq_col: (Optional) the column containing the reference sequence of the codon + :param codon_alt_seq_col: (Optional) the column containing the alternative sequence of the codon + :param replace_current_protein_variants: whether to replace current protein variants for the microhaplotypes in the table + :return: a reference to the updated PMO + """ + codon_location_cols = [codon_chrom_col, codon_start_col, codon_end_col] + if any(codon_location_cols) and not all(codon_location_cols): + raise ValueError( + "If any of codon_chrom_col, codon_start_col or codon_end_col are set, all three must be." + ) + PMOUpdater._check_columns( + protein_variants_info, + [ + target_name_col, + seq_col, + transcript_col, + protein_start_col, + protein_end_col, + protein_genome_id_col, + protein_ref_seq_col, + protein_alt_seq_col, + gene_name_col, + alternative_gene_name_col, + *codon_location_cols, + codon_genome_id_col, + codon_strand_col, + codon_ref_seq_col, + codon_alt_seq_col, + ], + "protein_variants_info", + ) + target_indexes = ( + PMOProcessor.get_index_key_of_target_in_representative_microhaplotypes(pmo) + ) + microhaplotypes = [ + PMOUpdater._get_representative_microhaplotype( + pmo, str(row[target_name_col]), str(row[seq_col]), target_indexes + ) + for _, row in protein_variants_info.iterrows() + ] + + for microhaplotype in microhaplotypes: + if replace_current_protein_variants or not microhaplotype.get( + "associated_protein_variants" + ): + microhaplotype["associated_protein_variants"] = [] + + for (_, row), microhaplotype in zip( + protein_variants_info.iterrows(), microhaplotypes + ): + context = f"protein variant for target '{row[target_name_col]}'" + variant = { + "protein_location": PMOUpdater._genomic_location_from_row( + row, + context, + protein_genome_id_col, + transcript_col, + protein_start_col, + protein_end_col, + ref_seq_col=protein_ref_seq_col, + alt_seq_col=protein_alt_seq_col, + ) + } + for key, col in ( + ("gene_name", gene_name_col), + ("alternative_gene_name", alternative_gene_name_col), + ): + if col and not PMOUpdater._is_blank(row[col]): + variant[key] = str(row[col]) + if codon_chrom_col and not all( + PMOUpdater._is_blank(row[col]) for col in codon_location_cols + ): + variant[ + "codon_genomic_location" + ] = PMOUpdater._genomic_location_from_row( + row, + f"codon of {context}", + codon_genome_id_col, + codon_chrom_col, + codon_start_col, + codon_end_col, + codon_strand_col, + codon_ref_seq_col, + codon_alt_seq_col, + ) + microhaplotype["associated_protein_variants"].append(variant) + return pmo + @staticmethod def merge_dicts_by_key( main_list: list[dict], diff --git a/tests/test_pmo_builder/test_pmo_updater.py b/tests/test_pmo_builder/test_pmo_updater.py index 930fa28..447474c 100755 --- a/tests/test_pmo_builder/test_pmo_updater.py +++ b/tests/test_pmo_builder/test_pmo_updater.py @@ -152,11 +152,11 @@ def test_update_specimen_meta_with_traveler_info_with_optional(self): "travel_country": "Kenya", "travel_start_date": "2024-01", "travel_end_date": "2024-01-20", - "bed_net": 0.5, - "admin1": "Nairobi", - "admin2": "SubCounty1", - "admin3": "Ward1", - "latlon": "-1.2921,36.8219", + "bed_net_usage": 0.5, + "geo_admin1": "Nairobi", + "geo_admin2": "SubCounty1", + "geo_admin3": "Ward1", + "lat_lon": "-1.2921,36.8219", } ], }, @@ -167,11 +167,11 @@ def test_update_specimen_meta_with_traveler_info_with_optional(self): "travel_country": "Tanzania", "travel_start_date": "2024-02", "travel_end_date": "2024-02-15", - "bed_net": 0.0, - "admin1": "Dar es Salaam", - "admin2": "SubCounty2", - "admin3": "Ward2", - "latlon": "-6.7924,39.2083", + "bed_net_usage": 0.0, + "geo_admin1": "Dar es Salaam", + "geo_admin2": "SubCounty2", + "geo_admin3": "Ward2", + "lat_lon": "-6.7924,39.2083", } ], }, @@ -233,11 +233,11 @@ def test_update_specimen_meta_with_traveler_info_with_optional_replace_old(self) "travel_country": "Kenya", "travel_start_date": "2024-01", "travel_end_date": "2024-01-20", - "bed_net": 0.5, - "admin1": "Nairobi", - "admin2": "SubCounty1", - "admin3": "Ward1", - "latlon": "-1.2921,36.8219", + "bed_net_usage": 0.5, + "geo_admin1": "Nairobi", + "geo_admin2": "SubCounty1", + "geo_admin3": "Ward1", + "lat_lon": "-1.2921,36.8219", } ], }, @@ -248,11 +248,11 @@ def test_update_specimen_meta_with_traveler_info_with_optional_replace_old(self) "travel_country": "Tanzania", "travel_start_date": "2024-02", "travel_end_date": "2024-02-15", - "bed_net": 0.0, - "admin1": "Dar es Salaam", - "admin2": "SubCounty2", - "admin3": "Ward2", - "latlon": "-6.7924,39.2083", + "bed_net_usage": 0.0, + "geo_admin1": "Dar es Salaam", + "geo_admin2": "SubCounty2", + "geo_admin3": "Ward2", + "lat_lon": "-6.7924,39.2083", } ], }, @@ -260,6 +260,273 @@ def test_update_specimen_meta_with_traveler_info_with_optional_replace_old(self) } self.assertEqual(test_out_pmo, test_pmo) + def test_update_specimen_meta_with_traveler_info_skips_blank_optional(self): + test_pmo = {"specimen_info": [{"specimen_name": "spec1"}]} + traveler_info = pd.DataFrame( + { + "specimen_name": ["spec1"], + "travel_country": ["Kenya"], + "travel_start_date": ["2024-01"], + "travel_end_date": ["2024-02"], + "bed_net": [None], + "admin1": [""], + } + ) + PMOUpdater.update_specimen_meta_with_traveler_info( + test_pmo, + traveler_info, + bed_net_usage_col="bed_net", + geo_admin1_col="admin1", + ) + self.assertEqual( + test_pmo["specimen_info"][0]["travel_out_six_month"], + [ + { + "travel_country": "Kenya", + "travel_start_date": "2024-01", + "travel_end_date": "2024-02", + } + ], + ) + + def test_update_target_info_with_markers_of_interest(self): + test_pmo = { + "target_info": [{"target_name": "t1"}, {"target_name": "t2"}], + } + markers_info = pd.DataFrame( + { + "target_name": ["t1", "t1"], + "chrom": ["Pf3D7_04_v3", "Pf3D7_04_v3"], + "start": [748238, 748409], + "end": [748239, 748410], + "ref_seq": ["A", ""], + "associations": ["SP resistance, dhfr", None], + } + ) + PMOUpdater.update_target_info_with_markers_of_interest( + test_pmo, + markers_info, + ref_seq_col="ref_seq", + associations_col="associations", + ) + self.assertEqual( + test_pmo["target_info"][0]["markers_of_interest"], + [ + { + "marker_location": { + "genome_id": 0, + "chrom": "Pf3D7_04_v3", + "start": 748238, + "end": 748239, + "ref_seq": "A", + }, + "associations": ["SP resistance", "dhfr"], + }, + { + "marker_location": { + "genome_id": 0, + "chrom": "Pf3D7_04_v3", + "start": 748409, + "end": 748410, + } + }, + ], + ) + self.assertNotIn("markers_of_interest", test_pmo["target_info"][1]) + + def test_update_target_info_with_markers_of_interest_replace(self): + old_marker = { + "marker_location": {"genome_id": 0, "chrom": "c", "start": 1, "end": 2} + } + markers_info = pd.DataFrame( + {"target_name": ["t1"], "chrom": ["c"], "start": [5], "end": [6]} + ) + for replace, expected_count in ((False, 2), (True, 1)): + test_pmo = { + "target_info": [ + {"target_name": "t1", "markers_of_interest": [dict(old_marker)]} + ] + } + PMOUpdater.update_target_info_with_markers_of_interest( + test_pmo, markers_info, replace_current_markers=replace + ) + self.assertEqual( + len(test_pmo["target_info"][0]["markers_of_interest"]), + expected_count, + ) + + def test_update_target_info_with_markers_of_interest_raises(self): + test_pmo = {"target_info": [{"target_name": "t1"}]} + unknown_target = pd.DataFrame( + {"target_name": ["t9"], "chrom": ["c"], "start": [1], "end": [2]} + ) + with self.assertRaises(ValueError): + PMOUpdater.update_target_info_with_markers_of_interest( + test_pmo, unknown_target + ) + blank_start = pd.DataFrame( + {"target_name": ["t1"], "chrom": ["c"], "start": [None], "end": [2]} + ) + with self.assertRaises(ValueError): + PMOUpdater.update_target_info_with_markers_of_interest( + test_pmo, blank_start + ) + with self.assertRaises(ValueError): + PMOUpdater.update_target_info_with_markers_of_interest( + test_pmo, unknown_target, chrom_col="not_a_column" + ) + + def _rep_mhap_pmo(self): + return { + "target_info": [{"target_name": "t1"}, {"target_name": "t2"}], + "representative_microhaplotypes": { + "targets": [ + { + "target_id": 1, + "microhaplotypes": [{"seq": "ACGT"}, {"seq": "ACTT"}], + } + ] + }, + } + + def test_update_representative_microhaplotypes_with_seq_variants(self): + test_pmo = self._rep_mhap_pmo() + seq_variants_info = pd.DataFrame( + { + "target_name": ["t2", "t2"], + "seq": ["ACTT", "ACTT"], + "chrom": ["c", "c"], + "start": [2, 3], + "end": [3, 4], + "ref_seq": ["G", "T"], + "alt_seq": ["T", None], + } + ) + PMOUpdater.update_representative_microhaplotypes_with_seq_variants( + test_pmo, + seq_variants_info, + ref_seq_col="ref_seq", + alt_seq_col="alt_seq", + ) + mhaps = test_pmo["representative_microhaplotypes"]["targets"][0][ + "microhaplotypes" + ] + self.assertNotIn("associated_seq_variants", mhaps[0]) + self.assertEqual( + mhaps[1]["associated_seq_variants"], + [ + { + "genome_id": 0, + "chrom": "c", + "start": 2, + "end": 3, + "ref_seq": "G", + "alt_seq": "T", + }, + {"genome_id": 0, "chrom": "c", "start": 3, "end": 4, "ref_seq": "T"}, + ], + ) + + def test_update_representative_microhaplotypes_with_seq_variants_raises(self): + for target_name, seq in (("t2", "GGGG"), ("t1", "ACGT")): + seq_variants_info = pd.DataFrame( + { + "target_name": [target_name], + "seq": [seq], + "chrom": ["c"], + "start": [1], + "end": [2], + } + ) + with self.assertRaises(ValueError): + PMOUpdater.update_representative_microhaplotypes_with_seq_variants( + self._rep_mhap_pmo(), seq_variants_info + ) + + def test_update_representative_microhaplotypes_with_protein_variants(self): + test_pmo = self._rep_mhap_pmo() + protein_variants_info = pd.DataFrame( + { + "target_name": ["t2", "t2"], + "seq": ["ACGT", "ACGT"], + "transcript": ["PF3D7_0417200.1", "PF3D7_0417200.1"], + "protein_start": [50, 107], + "protein_end": [51, 108], + "protein_ref": ["C", "S"], + "protein_alt": ["R", "N"], + "gene": ["dhfr", ""], + "codon_chrom": ["Pf3D7_04_v3", None], + "codon_start": [748238, None], + "codon_end": [748241, None], + } + ) + PMOUpdater.update_representative_microhaplotypes_with_protein_variants( + test_pmo, + protein_variants_info, + protein_ref_seq_col="protein_ref", + protein_alt_seq_col="protein_alt", + gene_name_col="gene", + codon_chrom_col="codon_chrom", + codon_start_col="codon_start", + codon_end_col="codon_end", + ) + mhaps = test_pmo["representative_microhaplotypes"]["targets"][0][ + "microhaplotypes" + ] + self.assertEqual( + mhaps[0]["associated_protein_variants"], + [ + { + "protein_location": { + "genome_id": 0, + "chrom": "PF3D7_0417200.1", + "start": 50, + "end": 51, + "ref_seq": "C", + "alt_seq": "R", + }, + "gene_name": "dhfr", + "codon_genomic_location": { + "genome_id": 0, + "chrom": "Pf3D7_04_v3", + "start": 748238, + "end": 748241, + }, + }, + { + "protein_location": { + "genome_id": 0, + "chrom": "PF3D7_0417200.1", + "start": 107, + "end": 108, + "ref_seq": "S", + "alt_seq": "N", + }, + }, + ], + ) + self.assertNotIn("associated_protein_variants", mhaps[1]) + + def test_update_representative_microhaplotypes_with_protein_variants_partial_codon_cols( + self, + ): + protein_variants_info = pd.DataFrame( + { + "target_name": ["t2"], + "seq": ["ACGT"], + "transcript": ["tx"], + "protein_start": [1], + "protein_end": [2], + "codon_chrom": ["c"], + } + ) + with self.assertRaises(ValueError): + PMOUpdater.update_representative_microhaplotypes_with_protein_variants( + self._rep_mhap_pmo(), + protein_variants_info, + codon_chrom_col="codon_chrom", + ) + # PMOUpdater.merge_dicts_by_key def test_merge_dicts_by_key_correct_fields_added(self): result = PMOUpdater.merge_dicts_by_key(