HLAfreq.HLAfreq

  1"""
  2Download and combine HLA allele frequencies from multiple datasets.
  3
  4Download allele frequency data from
  5[allelefrequencies.net](www.allelefrequencies.net). Allele
  6frequencies from different populations can be combined to
  7estimate HLA frequencies of countries or other regions such as
  8global HLA frequencies.
  9"""
 10
 11from collections.abc import Iterable
 12from bs4 import BeautifulSoup
 13import requests
 14import pandas as pd
 15import numpy as np
 16import matplotlib.pyplot as plt
 17import math
 18import scipy as sp
 19import matplotlib.colors as mcolors
 20
 21
 22def simulate_population(alleles: Iterable[str], locus: str, population: str):
 23    pop_size = np.random.randint(len(alleles), 50)
 24    samples = np.random.choice(alleles, pop_size, replace=True)
 25    counts = pd.Series(samples).value_counts()
 26    counts.values / pop_size
 27    pop = pd.DataFrame(
 28        {
 29            "allele": counts.index,
 30            "loci": locus,
 31            "population": population,
 32            "allele_freq": counts.values / pop_size,
 33            "sample_size": pop_size,
 34        }
 35    )
 36    return pop
 37
 38
 39def simulate_study(alleles, populations, locus):
 40    study = []
 41    for i in range(populations):
 42        pop = simulate_population(alleles=alleles, locus=locus, population=f"pop_{i}")
 43        study.append(pop)
 44
 45    study = pd.concat(study)
 46    return study
 47
 48
 49def makeURL(
 50    country="",
 51    standard="s",
 52    locus="",
 53    resolution_pattern="bigger_equal_than",
 54    resolution=2,
 55    region="",
 56    ethnic="",
 57    study_type="",
 58    dataset_source="",
 59    sample_year="",
 60    sample_year_pattern="",
 61    sample_size="",
 62    sample_size_pattern="",
 63):
 64    """Create URL for search of allele frequency net database.
 65
 66    All arguments are documented [here](https://www.allelefrequencies.net/extaccess.asp)
 67
 68    Args:
 69        country (str, optional): Country name to retrieve records from. Defaults to "".
 70        standard (str, optional): Filter study quality standard to this or higher.
 71            {'g', 's', 'a'} Gold, silver, all. Defaults to 's'.
 72        locus (str, optional): The locus to return allele data for. Defaults to "".
 73        resolution_pattern (str, optional): Resolution comparitor {'equal', 'different',
 74            'less_than', 'bigger_than', 'less_equal_than', 'bigger_equal_than'}.
 75            Filter created using `resolution` and `resolution_pattern`.
 76            Defaults to "bigger_equal_than".
 77        resolution (int, optional): Number of fields of resolution of allele. Filter
 78            created using `resolution` and `resolution_pattern`. Defaults to 2.
 79        region (str, optional): Filter to geographic region. {Asia, Australia,
 80            Eastern Europe, ...}.
 81            All regions listed [here](https://www.allelefrequencies.net/pop6003a.asp).
 82            Defaults to "".
 83        ethnic (str, optional): Filter to ethnicity. {"Amerindian", "Black", "Caucasian", ...}.
 84            All ethnicities listed [here](https://www.allelefrequencies.net/pop6003a.asp).
 85            Defaults to "".
 86        study_type (str, optional): Type of study. {"Anthropology", "Blood+Donor",
 87            "Bone+Marrow+Registry", "Controls+for+Disease+Study", "Disease+Study+Patients",
 88            "Other", "Solid+Organd+Unrelated+Donors", "Stem+cell+donors"}. Defaults to "".
 89        dataset_source (str, optional): Source of data. {"Literature",
 90            "Proceedings+of+IHWs", "Unpublished"}. Defaults to "".
 91        sample_year (int, optional): Sample year to compare to. Filter created using
 92            sample_year and sample_year_pattern. Defaults to "".
 93        sample_year_pattern (str, optional): Pattern to compare sample year to. Filter
 94            created using sample_year and sample_year_pattern. {'equal', 'different',
 95            'less_than', 'bigger_than', 'less_equal_than', 'bigger_equal_than'}.
 96            Defaults to "".
 97        sample_size (int, optional): Sample size to compare to. Filter created using
 98            sample_size and sample_size_pattern. Defaults to "".
 99        sample_size_pattern (str, optional): Pattern to compare sample size to. Filter
100            created using sample_size and sample_size_pattern. {'equal', 'different',
101            'less_than', 'bigger_than', 'less_equal_than', 'bigger_equal_than'}.
102            Defaults to "".
103
104    Returns:
105        str: URL to search allelefrequencies.net
106    """
107    base = "https://www.allelefrequencies.net/hla6006a.asp?"
108    locus_type = "hla_locus_type=Classical&"
109    hla_locus = "hla_locus=%s&" % (locus)
110    country = "hla_country=%s&" % (country)
111    region = "hla_region=%s&" % (region)
112    ethnic = "hla_ethnic=%s&" % (ethnic)
113    study_type = "hla_study=%s&" % (study_type)
114    dataset_source = "hla_dataset_source=%s&" % (dataset_source)
115    sample_year = "hla_sample_year=%s&" % (sample_year)
116    sample_year_pattern = "hla_sample_year_pattern=%s&" % (sample_year_pattern)
117    sample_size = "hla_sample_size=%s&" % (sample_size)
118    sample_size_pattern = "hla_sample_size_pattern=%s&" % (sample_size_pattern)
119    hla_level_pattern = "hla_level_pattern=%s&" % (resolution_pattern)
120    hla_level = "hla_level=%s&" % (resolution)
121    standard = "standard=%s&" % standard
122    url = (
123        base
124        + locus_type
125        + hla_locus
126        + country
127        + hla_level_pattern
128        + hla_level
129        + standard
130        + region
131        + ethnic
132        + study_type
133        + dataset_source
134        + sample_year
135        + sample_year_pattern
136        + sample_size
137        + sample_size_pattern
138    )
139    return url
140
141
142def parseAF(bs):
143    """Generate a dataframe from a given html page
144
145    Args:
146        bs (bs4.BeautifulSoup): BeautifulSoup object from allelefrequencies.net page
147
148    Returns:
149        pd.DataFrame: Table of allele, allele frequency, samplesize, and population
150    """
151    # Get the results table from the div `divGenDetail`
152    tab = bs.find("div", {"id": "divGenDetail"}).find("table", {"class": "tblNormal"})
153    # Get the column headers from the first row of the table
154    columns = [
155        "line",
156        "allele",
157        "flag",
158        "population",
159        "carriers%",
160        "allele_freq",
161        "AF_graphic",
162        "sample_size",
163        "database",
164        "distribution",
165        "haplotype_association",
166        "notes",
167    ]
168    rows = []
169    for row in tab.find_all("tr"):
170        rows.append([td.get_text(strip=True) for td in row.find_all("td")])
171    # Make dataframe of table rows
172    # skip the first row as it's `th` headers
173    df = pd.DataFrame(rows[1:], columns=columns)
174
175    # Get HLA loci
176    df["loci"] = df.allele.apply(lambda x: x.split("*")[0])
177
178    # Drop unwanted columns
179    df = df[["allele", "loci", "population", "allele_freq", "carriers%", "sample_size"]]
180    return df
181
182
183def Npages(bs):
184    """How many pages of results are there?
185
186    Args:
187        bs (bs4.BeautifulSoup): BS object of allelefrequencies.net results page
188
189    Returns:
190        int: Total number of results pages
191    """
192    # Get the table with number of pages
193    navtab = bs.find("div", {"id": "divGenNavig"}).find("table", {"class": "table10"})
194    if not navtab:
195        raise AssertionError(
196            "navtab does not evaluate to True. Check URL returns results in web browser."
197        )
198    # Get cell with ' of ' in
199    pagesOfN = [
200        td.get_text(strip=True) for td in navtab.find_all("td") if " of " in td.text
201    ]
202    # Check single cell returned
203    if not len(pagesOfN) == 1:
204        raise AssertionError("divGenNavig should contain 1 of not %s" % len(pagesOfN))
205    # Get total number of pages
206    N = pagesOfN[0].split("of ")[1]
207    N = int(N)
208    return N
209
210
211def formatAF(AFtab, ignoreG=True):
212    """Format allele frequency table.
213
214    Convert sample_size and allele_freq to numeric data type.
215    Removes commas from sample size. Removes "(*)" from allele frequency if
216    `ignoreG` is `True`. `formatAF()` is used internally by combineAF and getAFdata
217    by default.
218
219    Args:
220        AFtab (pd.DataFrame): Allele frequency data downloaded from allelefrequency.net
221            using `getAFdata()`.
222        ignoreG (bool, optional): Treat G group alleles as normal.
223            See https://hla.alleles.org/alleles/g_groups.html for details. Defaults to True.
224
225    Returns:
226        pd.DataFrame: The formatted allele frequency data.
227    """
228    df = AFtab.copy()
229    if pd.api.types.is_string_dtype(df.sample_size.dtype):
230        df.sample_size = pd.to_numeric(df.sample_size.str.replace(",", ""))
231    if pd.api.types.is_string_dtype(df.allele_freq.dtype):
232        if ignoreG:
233            df.allele_freq = df.allele_freq.str.replace("(*)", "", regex=False)
234        df.allele_freq = pd.to_numeric(df.allele_freq)
235    return df
236
237
238def getAFdata(base_url, timeout=20, format=True, ignoreG=True):
239    """Get all allele frequency data from a search base_url.
240
241    Iterates over all pages regardless of which page is based.
242
243    Args:
244        base_url (str): URL for base search.
245        timeout (int): How long to wait to receive a response.
246        format (bool): Format the downloaded data using `formatAF()`.
247        ignoreG (bool): treat allele G groups as normal.
248            See https://hla.alleles.org/alleles/g_groups.html for details. Default = True
249
250    Returns:
251        pd.DataFrame: allele frequency data parsed into a pandas dataframe
252    """
253    # Get BS object from base search
254    try:
255        bs = BeautifulSoup(requests.get(base_url, timeout=timeout).text, "html.parser")
256    except requests.exceptions.ReadTimeout as e:
257        raise Exception(
258            "Requests timeout, try a larger `timeout` value for `getAFdata()`"
259        ) from None
260    # How many pages of results
261    N = Npages(bs)
262    print("%s pages of results" % N)
263    # iterate over pages, parse and combine data from each
264    tabs = []
265    for i in range(N):
266        # print (" Parsing page %s" %(i+1))
267        print(" Parsing page %s" % (i + 1), end="\r")
268        url = base_url + "page=" + str(i + 1)
269        try:
270            bs = BeautifulSoup(requests.get(url, timeout=timeout).text, "html.parser")
271        except requests.exceptions.ReadTimeout as e:
272            raise Exception(
273                "Requests timeout, try a larger `timeout` value for `getAFdata()`"
274            ) from None
275        tab = parseAF(bs)
276        tabs.append(tab)
277    print("Download complete")
278    tabs = pd.concat(tabs)
279    if format:
280        try:
281            tabs = formatAF(tabs, ignoreG)
282        except AttributeError:
283            print("Formatting failed, non-numeric datatypes may remain.")
284    return tabs
285
286
287def incomplete_studies(AFtab, llimit=0.95, ulimit=1.1, datasetID="population"):
288    """Report any studies with allele freqs that don't sum to 1
289
290    Args:
291        AFtab (pd.DataFrame): Dataframe containing multiple studies
292        llimit (float, optional): Lower allele_freq sum limit that counts as complete.
293            Defaults to 0.95.
294        ulimit (float, optional): Upper allele_freq sum limit that will not be reported.
295            Defaults to 1.1.
296        datasetID (str): Unique identifier column for study
297    """
298    poplocs = AFtab.groupby([datasetID, "loci"]).allele_freq.sum()
299    lmask = poplocs < llimit
300    if sum(lmask > 0):
301        print(poplocs[lmask])
302        print(f"{sum(lmask)} studies have total allele frequency < {llimit}")
303    umask = poplocs > ulimit
304    if sum(umask > 0):
305        print(poplocs[umask])
306        print(f"{sum(umask)} studies have total allele frequency > {ulimit}")
307    incomplete = pd.concat([poplocs[lmask], poplocs[umask]])
308    return incomplete
309
310
311def only_complete(AFtab, llimit=0.95, ulimit=1.1, datasetID="population"):
312    """Returns only complete studies.
313
314    Data is dropped if the locus for that population is not complete, i.e. doesn't
315    sum to between `llimit` and `ulimit`. This prevents throwing away data if
316    another loci in the population is incomplete.
317
318    Args:
319        AFtab (pd.DataFrame): Dataframe containing multiple studies
320        llimit (float, optional): Lower allele_freq sum limit that counts as complete.
321            Defaults to 0.95.
322        ulimit (float, optional): Upper allele_freq sum limit that will not be reported.
323            Defaults to 1.1.
324        datasetID (str): Unique identifier column for study. Defaults to 'population'.
325
326    Returns:
327        pd.DataFrame: Allele frequency data of multiple studies, but only complete studies.
328    """
329    noncomplete = incomplete_studies(
330        AFtab=AFtab, llimit=llimit, ulimit=ulimit, datasetID=datasetID
331    )
332    # Returns False if population AND loci are in the noncomplete.index
333    # AS A PAIR
334    # This is important so that we don't throw away all data on a population
335    # just because one loci is incomplete.
336    complete_mask = AFtab.apply(
337        lambda x: (x[datasetID], x.loci) not in noncomplete.index, axis=1
338    )
339    df = AFtab[complete_mask]
340    return df
341
342
343def check_resolution(AFtab):
344    """Check if all alleles in AFtab have the same resolution.
345    Will print the number of records with each resolution.
346
347    Args:
348        AFtab (pd.DataFrame): Allele frequency data
349
350    Returns:
351        bool: True only if all alleles have the same resolution, else False.
352    """
353    resolution = 1 + AFtab.allele.str.count(":")
354    resVC = resolution.value_counts()
355    pass_check = len(resVC) == 1
356    if not pass_check:
357        print(resVC)
358        print("Multiple resolutions in AFtab. Fix with decrease_resolution()")
359    return pass_check
360
361
362def decrease_resolution(AFtab, newres, datasetID="population"):
363    """Decrease allele resolution so all alleles have the same resolution.
364
365    Args:
366        AFtab (pd.DataFrame): Allele frequency data.
367        newres (int): The desired number of fields for resolution.
368        datasetID (str, optional): Column to use as stud identifier.
369            Defaults to 'population'.
370
371    Returns:
372        pd.DataFrame: Allele frequency data with all alleles of requested resolution.
373    """
374    df = AFtab.copy()
375    resolution = 1 + df.allele.str.count(":")
376    if not all(resolution >= newres):
377        raise AssertionError(f"Some alleles have resolution below {newres} fields")
378    new_allele = df.allele.str.split(":").apply(lambda x: ":".join(x[:newres]))
379    df.allele = new_allele
380    collapsed = collapse_reduced_alleles(df, datasetID=datasetID)
381    return collapsed
382
383
384def collapse_reduced_alleles(AFtab, datasetID="population"):
385    df = AFtab.copy()
386    # Group by alleles within datasets
387    grouped = df.groupby([datasetID, "allele"])
388    # Sum allele freq but keep other columns
389    collapsed = grouped.apply(
390        lambda row: [
391            sum(row.allele_freq),
392            row.sample_size.unique()[0],
393            row.loci.unique()[0],
394            len(row.loci.unique()),
395            len(row.sample_size.unique()),
396        ],
397        include_groups=False
398    )
399    collapsed = pd.DataFrame(
400        collapsed.tolist(),
401        index=collapsed.index,
402        columns=["allele_freq", "sample_size", "loci", "#loci", "#sample_sizes"],
403    ).reset_index()
404    # Within a study each all identical alleles should have the same loci and sample size
405    if not all(collapsed["#loci"] == 1):
406        raise AssertionError(
407            "Multiple loci found for a single allele in a single population"
408        )
409    if not all(collapsed["#sample_sizes"] == 1):
410        raise AssertionError(
411            "Multiple sample_sizes found for a single allele in a single population"
412        )
413    collapsed = collapsed[
414        ["allele", "loci", "population", "allele_freq", "sample_size"]
415    ]
416    alleles_unique_in_study(collapsed)
417    return collapsed
418
419
420def unmeasured_alleles(AFtab, datasetID="population"):
421    """When combining AF estimates, unreported alleles can inflate frequencies
422        so AF sums to >1. Therefore we add unreported alleles with frequency zero.
423
424    Args:
425        AFtab (pd.DataFrame): Formatted allele frequency data
426        datasetID (str): Unique identifier column for study
427
428    Returns:
429        pd.DataFrame: Allele frequency data with all locus alleles reported
430            for each dataset
431    """
432    df = AFtab.copy()
433    loci = df.loci.unique()
434    # Iterate over loci separately
435    for locus in loci:
436        # Iterate over each dataset reporting that locus
437        datasets = df[df.loci == locus][datasetID].unique()
438        for dataset in datasets:
439            # Single locus, single dataset
440            datasetAF = df[(df[datasetID] == dataset) & (df.loci == locus)]
441            # What was the sample size for this data?
442            dataset_sample_size = datasetAF.sample_size.unique()
443            if not (len(dataset_sample_size) == 1):
444                raise AssertionError(
445                    "dataset_sample_size must be 1, not %s" % len(dataset_sample_size)
446                )
447            dataset_sample_size = dataset_sample_size[0]
448            # Get all alleles for this locus (across datasets)
449            ualleles = df[df.loci == locus].allele.unique()
450            # Which of these alleles are not in this dataset?
451            missing_alleles = [
452                allele for allele in ualleles if not allele in datasetAF.allele.values
453            ]
454            missing_rows = [
455                (al, locus, dataset, 0, 0, dataset_sample_size)
456                for al in missing_alleles
457            ]
458            missing_rows = pd.DataFrame(
459                missing_rows,
460                columns=[
461                    "allele",
462                    "loci",
463                    datasetID,
464                    "allele_freq",
465                    "carriers%",
466                    "sample_size",
467                ],
468            )
469            # Add them in with zero frequency
470            if not missing_rows.empty:
471                df = pd.concat([df, missing_rows], ignore_index=True)
472    return df
473
474
475def combineAF(
476    AFtab,
477    weights="2n",
478    alpha=[],
479    datasetID="population",
480    format=True,
481    ignoreG=True,
482    add_unmeasured=True,
483    complete=True,
484    resolution=True,
485    unique=True,
486):
487    """Combine allele frequencies from multiple studies.
488
489    `datasetID` is the unique identifier for studies to combine.
490    Allele frequencies combined using a Dirichlet distribution where each study's
491    contribution to the concentration parameter is $2 * sample_size * allele_frequency$.
492    Sample size is doubled to get `2n` due to diploidy. If an alternative `weights` is
493    set it is not doubled. The total concentration parameter of the Dirichlet distribution
494    is the contributions from all studies plus the prior `alpha`. If `alpha` is not set
495    the prior defaults to 1 observation of each allele.
496
497    Args:
498        AFtab (pd.DataFrame): Table of Allele frequency data
499        weights (str, optional): Column to be weighted by allele frequency to generate
500            concentration parameter of Dirichlet distribution. Defaults to '2n'.
501        alpha (list, optional): Prior to use for Dirichlet distribution. Defaults to [].
502        datasetID (str, optional): Unique identifier column for study. Defaults to
503            'population'.
504        format (bool, optional): Run `formatAF()`. Defaults to True.
505        ignoreG (bool, optional): Treat allele G groups as normal, see `formatAF()`.
506            Defaults to True.
507        add_unmeasured (bool, optional): Add unmeasured alleles to each study. This is
508            important to ensure combined allele frequencies sum to 1. See
509            `add_unmeasured()`. Defaults to True.
510        complete (bool, optional): Check study completeness. Uses default values for
511            `incomplete_studies()`. If you are happy with your study completeness can
512            be switched off with False. Defaults to True.
513        resolution (bool, optional): Check that all alleles have the same resolution,
514            see `check_resolution()`. Defaults to True.
515        unique (bool, optional): Check that each allele appears no more than once per
516            study. See `alleles_unique_in_study()`. Defaults to True.
517
518    Returns:
519        pd.DataFrame: Allele frequencies after combining estimates from all studies.
520            *allele_freq* is the combined frequency estimate from the Dirichlet mean
521            where the concentration is `alpha` + `c`.
522            *alpha* is the prior used for the Dirichlet distribution.
523            *c* is the observations used for the Dirichlet distribution.
524            *sample_size* is the total sample size of all combined studies.
525            *wav* is the weighted average.
526    """
527    df = AFtab.copy()
528    single_loci(df)
529    if unique:
530        if not alleles_unique_in_study(df, datasetID=datasetID):
531            raise AssertionError("The same allele appears multiple times in a dataset")
532    if complete:
533        if not incomplete_studies(df, datasetID=datasetID).empty:
534            raise AssertionError(
535                "AFtab contains studies with AF that doesn't sum to 1. Check incomplete_studies(AFtab)"
536            )
537    if resolution:
538        if not check_resolution(df):
539            raise AssertionError(
540                "AFtab conains alleles at multiple resolutions, check check_resolution(AFtab)"
541            )
542    if format:
543        df = formatAF(df, ignoreG)
544    if add_unmeasured:
545        df = unmeasured_alleles(df, datasetID)
546    try:
547        df["2n"] = df.sample_size * 2
548    except AttributeError:
549        print("column '2n' could not be created")
550    df["c"] = df.allele_freq * df[weights]
551    grouped = df.groupby("allele", sort=True)
552    combined = grouped.apply(
553        lambda row: [
554            row.name,
555            row.loci.unique()[0],
556            np.average(row.allele_freq, weights=row[weights]),
557            row.c.sum(),
558            row.sample_size.sum(),
559        ],
560        include_groups=False
561    )
562    combined = pd.DataFrame(
563        combined.tolist(), columns=["allele", "loci", "wav", "c", "sample_size"]
564    )
565    combined = combined.reset_index(drop=True)
566    # Check that all alleles in a locus have the same sample size
567    # after merging
568    if duplicated_sample_size(combined):
569        id_duplicated_allele(grouped)
570    if not alpha:
571        alpha = default_prior(len(combined.allele))
572    combined["alpha"] = alpha
573    # Calculate Dirichlet mean for each allele
574    combined["allele_freq"] = sp.stats.dirichlet(combined.alpha + combined.c).mean()
575
576    return combined
577
578
579def default_prior(k):
580    """Calculate a default prior, 1 observation of each class.
581
582    Args:
583        k (int): Number of classes in the Dirichlet distribution.
584
585    Returns:
586        list: List of k 1s to use as prior.
587    """
588    alpha = [1] * k
589    return alpha
590
591
592def single_loci(AFtab):
593    """Check that allele frequency data is only of one locus
594
595    Args:
596        AFtab (pd.DataFrame): Allele frequency data
597    """
598    if not len(AFtab.loci.unique()) == 1:
599        raise AssertionError("'AFtab' must contain only 1 loci")
600
601
602def alleles_unique_in_study(AFtab, datasetID="population"):
603    """Are all alleles unique in each study?
604
605    Checks that no alleles are reported more than once in a single study.
606    Study is defined by `datasetID`.
607
608    Args:
609        AFtab (pd.DataFrame): Allele frequency data
610        datasetID (str, optional): Unique identifier column to define study.
611            Defaults to 'population'.
612
613    Returns:
614        bool: `True` on if no alleles occur more than once in any study, otherwise `False`.
615    """
616    df = AFtab.copy()
617    grouped = df.groupby([datasetID, "allele"])
618    # Are allele alleles unique? i.e. do any occur multiple times in grouping?
619    unique = grouped.size()[grouped.size() > 1].empty
620    if not unique:
621        print(f"Non unique alleles in study, is datasetID correct? {datasetID}")
622        print(grouped.size()[grouped.size() > 1])
623    return unique
624
625
626def duplicated_sample_size(AFtab):
627    """Returns True if any loci has more than 1 unique sample size"""
628    locus_sample_sizes = AFtab.groupby("loci").sample_size.apply(
629        lambda x: len(x.unique())
630    )
631    return any(locus_sample_sizes != 1)
632
633
634def id_duplicated_allele(grouped):
635    """Reports the allele that has mupltiple sample sizes"""
636    duplicated_population = grouped.population.apply(lambda x: any(x.duplicated()))
637    if not all(~duplicated_population):
638        raise AssertionError(
639            f"duplicated population within allele {duplicated_population[duplicated_population].index.tolist()}"
640        )
641
642
643def population_coverage(p):
644    """Proportion of people with at least 1 copy of this allele assuming HWE.
645
646    Args:
647        p (float): Allele frequency
648
649    Returns:
650        float: Sum of homozygotes and heterozygotes for this allele
651    """
652    q = 1 - p
653    homo = p**2
654    hetero = 2 * p * q
655    return homo + hetero
656
657
658def betaAB(alpha):
659    """Calculate `a` `b` values for all composite beta distributions.
660
661    Given the `alpha` vector defining a Dirichlet distribution calculate the `a` `b` values
662    for all composite beta distributions.
663
664    Args:
665        alpha (list): Values defining a Dirichlet distribution. This will be the prior
666            (for a naive distribution) or the prior + caf.c for a posterior distribution.
667
668    Returns:
669        list: List of `a` `b` values defining beta values, i.e. for each allele it is
670            the number of times it was and wasn't observed.
671    """
672    ab = [(a, sum(alpha) - a) for a in alpha]
673    return ab
674
675
676# def betaCI(a,b,credible_interval=0.95):
677#     """Calculat the central credible interval of a beta distribution
678
679#     Args:
680#         a (float): Beta shape parameter `a`, i.e. the number of times the allele was observed.
681#         b (float): Beta shape parameter `b`, i.e. the number of times the allele was not observed.
682#         credible_interval (float, optional): The size of the credible interval requested. Defaults to 0.95.
683
684#     Returns:
685#         tuple: Lower and upper credible interval of beta distribution.
686#     """
687#     bd = sp.stats.beta(a,b)
688#     lower_quantile = (1-credible_interval)/2
689#     upper_quantile = 1-lower_quantile
690#     lower_interval = bd.ppf(lower_quantile)
691#     upper_interval = bd.ppf(upper_quantile)
692#     return lower_interval, upper_interval
693
694# def AFci(caf, credible_interval=0.95):
695#     """Calculate credible interval for combined allele frequency table.
696#     Note that this ignores sampling error so confidence interval is too tight.
697#     Use HLAhdi.AFhdi() instead.
698
699#     Args:
700#         caf (pd.DataFrame): Table produced by combineAF()
701#         credible_interval (float, optional): The desired confidence interval. Defaults to 0.95.
702
703#     Returns:
704#         list: Lower and upper credible intervals as a list of tuples
705#     """
706#     ab = betaAB(
707#         caf.alpha + caf.c,
708#     )
709#     ci = [betaCI(a, b, credible_interval) for a,b in ab]
710#     return ci
711
712
713def plot_prior(concentration, ncol=2, psteps=1000, labels=""):
714    """Plot probability density function for prior values.
715
716    Args:
717        concentration (list): Vector of the prior Dirichlet concentration values.
718        ncol (int, optional): Number of columns. Defaults to 2.
719        labels (list, optional): Labels for elements of concentration in the same
720            order. Defaults to "".
721    """
722    ab = betaAB(concentration)
723    pline = np.linspace(0, 1, psteps)
724    nrow = math.ceil(len(concentration) / ncol)
725    fig, ax = plt.subplots(nrow, ncol, sharex=True)
726    fig.suptitle("Probability density")
727    # If labels is a list nothing happens,
728    # But if it's a series it converts to a list
729    labels = list(labels)
730    if not labels:
731        labels = [""] * len(concentration)
732    if not len(concentration) == len(labels):
733        raise AssertionError("concentration must be same length as labels")
734    for i, alpha in enumerate(concentration):
735        a, b = ab[i]
736        bd = sp.stats.beta(a, b)
737        pdf = [bd.pdf(p) for p in pline]
738        ax.flat[i].plot(pline, pdf)
739        ax.flat[i].set_title(labels[i])
740    for axi in ax[-1, :]:
741        axi.set(xlabel="Allele freq")
742    for axi in ax[:, 0]:
743        axi.set(ylabel="PDF")
744    plt.show()
745
746
747def plotAF(
748    caf=pd.DataFrame(),
749    AFtab=pd.DataFrame(),
750    cols=list(mcolors.TABLEAU_COLORS.keys()),
751    datasetID="population",
752    hdi=pd.DataFrame(),
753    compound_mean=pd.DataFrame(),
754):
755    """Plot allele frequency results from `HLAfreq`.
756
757    Plot combined allele frequencies, individual allele frequencies,
758    and credible intervals on combined allele frequency estimates.
759    Credible interval is only plotted if a value is given for `hdi`.
760    The plotted Credible interval is whatever was passed to HLAfreq_pymc.AFhdi()
761    when calculating hdi.
762
763    Args:
764        caf (pd.DataFrame, optional): Combined allele frequency estimates from
765            HLAfreq.combineAF. Defaults to pd.DataFrame().
766        AFtab (pd.DataFrame, optional): Table of allele frequency data. Defaults
767            to pd.DataFrame().
768        cols (list, optional): List of colours to use for each individual dataset.
769            Defaults to list(mcolors.TABLEAU_COLORS.keys()).
770        datasetID (str, optional): Column used to define separate datasets. Defaults
771            to "population".
772        weights (str, optional): Column to be weighted by allele frequency to generate
773            concentration parameter of Dirichlet distribution. Defaults to '2n'.
774        hdi (pd.DataFrame, optional): The high density interval object to plot credible
775            intervals. Produced by HLAfreq.HLA_pymc.AFhdi(). Defaults to pd.DataFrame().
776        compound_mean (pd.DataFrame, optional): The high density interval object to plot
777            post_mean. Produced by HLAfreq.HLA_pymc.AFhdi(). Defaults to pd.DataFrame().
778    """
779    # Plot allele frequency for each dataset
780    if not AFtab.empty:
781        # Cols must be longer than the list of alleles
782        # If not, repeat the list of cols
783        repeat_cols = np.ceil(len(AFtab[datasetID]) / len(cols))
784        repeat_cols = int(repeat_cols)
785        cols = cols * repeat_cols
786        # Make a dictionary mapping datasetID to colours
787        cmap = dict(zip(AFtab[datasetID].unique(), cols))
788        plt.scatter(
789            x=AFtab.allele_freq,
790            y=AFtab.allele,
791            c=[cmap[x] for x in AFtab[datasetID]],
792            alpha=0.7,
793            zorder=2,
794        )
795    # Plot combined allele frequency
796    if not caf.empty:
797        plt.scatter(
798            x=caf.allele_freq,
799            y=caf.allele,
800            edgecolors="black",
801            facecolors="none",
802            zorder=3,
803        )
804    # Plot high density interval
805    if not hdi.empty:
806        # assert not AFtab.empty, "AFtab is needed to calculate credible interval"
807        # from HLAfreq import HLAfreq_pymc as HLAhdi
808        # print("Fitting model with PyMC, make take a few seconds")
809        # hdi = HLAhdi.AFhdi(
810        #     AFtab=AFtab,
811        #     weights=weights,
812        #     datasetID=datasetID,
813        #     credible_interval=credible_interval,
814        #     conc_mu=conc_mu,
815        #     conc_sigma=conc_sigma
816        # )
817        for interval in hdi.iterrows():
818            # .iterrows returns a index and data as a tuple for each row
819            plt.hlines(
820                y=interval[1]["allele"],
821                xmin=interval[1]["lo"],
822                xmax=interval[1]["hi"],
823                color="black",
824            )
825    if not compound_mean.empty:
826        for row in compound_mean.iterrows():
827            plt.scatter(
828                y=row[1]["allele"], x=row[1]["post_mean"], color="black", marker="|"
829            )
830    plt.xlabel("Allele frequency")
831    plt.grid(zorder=0)
832    plt.show()