remove pre-py38 type annotations

This commit is contained in:
Stefan Steininger
2024-01-31 11:47:40 +01:00
parent 2e3fdb62e9
commit a5393d793e

View File

@@ -1,24 +1,27 @@
import math import math
from typing import Tuple, Dict
from functools import cache from functools import cache
import torch import torch
import torch.distributions as dist import torch.distributions as dist
from torch import exp from torch import exp
F1: Dict[float, float] = {} total_count = 0
F2: Dict[float, float] = {} calculation_count = 0
CDF: Dict[float, float] = {}
lookup_count = 0
def get_lookup_count(): def get_lookup_count():
global lookup_count global total_count, calculation_count
return lookup_count return total_count - calculation_count
@cache @cache
def calc_cdf(alpha: float) -> Tuple[float, float, float]: def calc_cdf(alpha: float) -> tuple[float, float, float]:
"""
Returns the calculated CDF and parameters f1,f2 from the input alpha
"""
global calculation_count
calculation_count += 1
normal = dist.Normal(0, 1) normal = dist.Normal(0, 1)
cdf_alpha = normal.cdf(torch.tensor(alpha)).item() cdf_alpha = normal.cdf(torch.tensor(alpha)).item()
pdf_alpha = exp(normal.log_prob(torch.tensor(alpha))).item() pdf_alpha = exp(normal.log_prob(torch.tensor(alpha))).item()
@@ -27,12 +30,8 @@ def calc_cdf(alpha: float) -> Tuple[float, float, float]:
1 - 2 * cdf_alpha) * alpha * pdf_alpha - pdf_alpha ** 2 1 - 2 * cdf_alpha) * alpha * pdf_alpha - pdf_alpha ** 2
return cdf_alpha, f1, f2 return cdf_alpha, f1, f2
def max_gaussian(mu1, sigma1, mu2, sigma2) -> Tuple[float, float]:
global lookup_count
global F1
global F2
global CDF
def max_gaussian(mu1, sigma1, mu2, sigma2) -> tuple[float, float]:
""" """
Returns the combined max gaussian of two Gaussians represented by mu1, sigma1, mu2, simga2 Returns the combined max gaussian of two Gaussians represented by mu1, sigma1, mu2, simga2
:param mu1: mu of the first Gaussian :param mu1: mu of the first Gaussian
@@ -41,10 +40,15 @@ def max_gaussian(mu1, sigma1, mu2, sigma2) -> Tuple[float, float]:
:param sigma2: sigma of the second Gaussian :param sigma2: sigma of the second Gaussian
:return: mu and sigma maximized :return: mu and sigma maximized
""" """
global total_count
total_count += 1
# we assume independence of the two gaussians # we assume independence of the two gaussians
# print(mu1, sigma1, mu2, sigma2) # print(mu1, sigma1, mu2, sigma2)
# normal = dist.Normal(0, 1) # normal = dist.Normal(0, 1)
sigma_m = math.sqrt(sigma1 ** 2 + sigma2 ** 2) sigma_m = math.sqrt(sigma1 ** 2 + sigma2 ** 2)
# round to two significant digits to enable float lookup
alpha = round((mu1 - mu2) / sigma_m, 2) alpha = round((mu1 - mu2) / sigma_m, 2)
cdf_alpha, f1_alpha, f2_alpha = calc_cdf(alpha) cdf_alpha, f1_alpha, f2_alpha = calc_cdf(alpha)
@@ -75,13 +79,13 @@ def min_gaussian(mu1, sigma1, mu2, sigma2) -> tuple[float, float]:
pdf_alpha_neg = exp(normal.log_prob(torch.tensor(-alpha))).item() pdf_alpha_neg = exp(normal.log_prob(torch.tensor(-alpha))).item()
mu = mu1 * (1 - cdf_alpha) + mu2 * cdf_alpha - pdf_alpha_neg * sigma_m mu = mu1 * (1 - cdf_alpha) + mu2 * cdf_alpha - pdf_alpha_neg * sigma_m
sigma = math.sqrt((mu1**2 + sigma1**2) * (1 - cdf_alpha) + (mu2**2 + sigma2**2) * cdf_alpha - (mu1 + mu2) * sigma_m * pdf_alpha - mu**2) sigma = math.sqrt((mu1 ** 2 + sigma1 ** 2) * (1 - cdf_alpha) + (mu2 ** 2 + sigma2 ** 2) * cdf_alpha - (
mu1 + mu2) * sigma_m * pdf_alpha - mu ** 2)
return mu, sigma return mu, sigma
except ValueError: except ValueError:
print(mu1, sigma1, mu2, sigma2) print(mu1, sigma1, mu2, sigma2)
def beta_mean(alpha, beta): def beta_mean(alpha, beta):
return alpha / (alpha + beta) return alpha / (alpha + beta)