text:
from math import ceil
from typing import Optional

from scipy.stats import t as t_dist

import numpy as np

def sample_size_bcra3f2(
    rho2: float,
    n: float,
    J: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed: bool = True,
    K0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    g2: int = 0,
    r21: float = 0,
    r22: float = 0,
    print_pretty: bool = True,
) -> Optional[int]:
    """Calculates the minimum sample size of a Three-Level Blocked (Fixed) Cluster-level Random Assignment Design,
    Treatment at Level 2

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    J: float
        Harmonic mean of level 2 units across level 3 units (or simple average)
    effect_size: float, default=0.25
        The effecet size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type I error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    K0: int, default=10
        Our initial guess for the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.5
        Average proportion of level 2 units randomly assigned to treatment within level 3 units
    g2: int, default=0
        Number of covariates at level 2
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r22: float, default=0
        Proportion of level 2 variance in the outcome explained by level 2 covariates
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = K0 * (J - 2) - g2
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        K1 = pow(M / effect_size, 2) * (
            rho2 * (1 - r22) / (p * (1 - p) * J)
            + (1 - rho2) * (1 - r21) / (p * (1 - p) * J * n)
        )
        if abs(K1 - K0) < tol:
            break
        K0 = (K1 + K0) / 2
    sample_size = ceil(K0) if df > 0 else None
    if print_pretty:
        print(f"K = {sample_size}")
    return sample_size

def sample_size_bcra3r2(
    rho2: float,
    rho3: float,
    omega3: float,
    n: float,
    J: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed: bool = True,
    K0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    g3: int = 0,
    r21: float = 0,
    r22: float = 0,
    r2t3: float = 0,
    print_pretty: bool = True,
) -> Optional[int]:
    """Calculates the power of a Three-Level Blocked Cluster-level Random Assignment Design,
    Treatment at Level 2

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    rho3: float
        Proportion of variance in the outcome between level 3 units (unconditional ICC3)
    omega3: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 3 units to the residual
        variance at level 3
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    J: float
        Harmonic mean of level 2 units across level 3 units (or simple average)
    effect_size: float, default=0.25
        Effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    K0: int, default=10
        Our initial guess for the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.5
        Average proportion of level 2 units randomly assigned to the treatment within level 3 units
    g3: int, default=0
        Number of covariates at level 3
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r22: float, default=0
        Proportion of level 2 variance in the outcome explained by level 2 covariates
    r2t3: float, default=0
        Proportion of treatment effect variance among level 3 units explained by level 3 covariates
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = K0 - g3 - 1
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        K1 = pow(M / effect_size, 2) * (
            rho3 * omega3 * (1 - r2t3)
            + rho2 * (1 - r22) / (p * (1 - p) * J)
            + (1 - rho3 - rho2) * (1 - r21) / (p * (1 - p) * J * n)
        )
        if abs(K1 - K0) < tol:
            break
        K0 = (K1 + K0) / 2
    sample_size = ceil(K0) if df > 0 else None
    if print_pretty:
        print(f"K = {sample_size}")
    return sample_size

def sample_size_bcra4f3(
    rho2: float,
    rho3: float,
    n: float,
    J: float,
    K: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.05,
    two_tailed: bool = True,
    L0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    g3: int = 0,
    r21: float = 0,
    r22: float = 0,
    r23: float = 0,
    print_pretty: bool = True,
) -> Optional[int]:
    """Calculates the minimum sample size of a Four-Level Blocked (Fixed) Cluster-level Random Assignment Design,
    Treatment at Level 3

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    rho3: float
        Proportion of variance in the outcome between level 3 units (unconditional ICC3)
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    J: float
        Harmonic mean of level 2 units across level 3 units (or simple average)
    K: float
        Harmonic mean of level 3 units across level 4 units (or simple average)
    effect_size: float, default=0.25
        Effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.1
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    L0: int, default=10
        Our initial guess of our minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.5
        Average proportion of level 3 units randomly assigned to the treatment within level 4 units
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r22: float, default=0
        Proportion of level 2 variance in the outcome explained by level 2 covariates
    r23: float, default=0
        Proportion of level 3 variance in the outcome explained by level 3 covariates
    g3: int, default=0
        Number of covariates at level 3
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = L0 * (K - 2) - g3
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        L1 = pow(M / effect_size, 2) * (
            rho3 * (1 - r23) / (p * (1 - p) * K)
            + rho2 * (1 - r22) / (p * (1 - p) * K * J)
            + (1 - rho3 - rho2) * (1 - r21) / (p * (1 - p) * J * K * n)
        )
        if abs(L1 - L0) < tol:
            break
        L0 = (L1 + L0) / 2
    sample_size = ceil(L0) if df > 0 else None
    if print_pretty:
        print(f"L = {sample_size}")
    return sample_size

def sample_size_bcra4r2(
    rho2: float,
    rho3: float,
    rho4: float,
    omega3: float,
    omega4: float,
    n: float,
    J: float,
    K: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed: bool = True,
    L0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    r21: float = 0,
    r22: float = 0,
    r2t3: float = 0,
    r2t4: float = 0,
    g4: int = 0,
    print_pretty: bool = True,
) -> Optional[int]:
    """Calculates the minimum sample size of a Four-Level Blocked Cluster-level Random Assignment Design, Treatment
    at Level 2

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    rho3: float
        Proportion of variance in the outcome between level 3 units (unconditional ICC3)
    rho4: float
        Proportion of variance in the outcome between level 4 untis (unconditional ICC4)
    omega3: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 3 units to the residual
        variance at level 3
    omega4: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 4 units to the residual
        variance at level 4
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    J: float
        Harmonic mean of level 2 units across level 3 units (or simple average)
    K: float
        Harmonic mean of level 3 units across level 4 units (or simple average)
    effect_size: float, default=0.25
        Effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    L0: int, default=10
        Our initial guess of our minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.5
        Average proportion of level 2 units randomly assignewd to treatment within level 3 units
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r22: float, default=0
        Proportion of level 2 variance in the outcome explained by level 2 covariates
    r2t3: float, default=0
        Proportion of treatment effect variance among level 3 units explained by level 3 covariates
    r2t4: float, default=0
        Proportion of treatment effect variance among level 4 units explained by level 4 covariates
    g4: int, default=0
        Number of covariates at level 4
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = L0 - g4 - 1
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        L1 = pow(M / effect_size, 2) * (
            rho4 * omega4 * (1 - r2t4)
            + rho3 * omega3 * (1 - r2t3) / K
            + rho2 * (1 - r22) / (p * (1 - p) * K * J)
            + (1 - rho4 - rho3 - rho2) * (1 - r21) / (p * (1 - p) * J * K * n)
        )
        if abs(L1 - L0) < tol:
            break
        L0 = (L1 + L0) / 2
    sample_size = ceil(L0) if df > 0 else None
    if print_pretty:
        print(f"L = {sample_size}")
    return sample_size

def sample_size_bcra4r3(
    rho2: float,
    rho3: float,
    rho4: float,
    omega4: float,
    n: float,
    J: float,
    K: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed: bool = True,
    L0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    r21: float = 0,
    r22: float = 0,
    r23: float = 0,
    r2t4: float = 0,
    g4: int = 0,
    print_pretty: bool = True,
) -> Optional[int]:
    """Calculates the minimum sample size of a Four-Level Blocked Cluster-level Random Assignment Design, Treatment
     at Level 3

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    rho3: float
        Proportion of variance in the outcome between level 3 units (unconditional ICC3)
    rho4: float
        Proportion of variance in the outcome between level 4 units (unconditional ICC4)
    omega4: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 4 units to the residual
        variance at level 4
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    J: float
        Harmonic mean of level 2 units across level 3 units (or simple average)
    K: float
        Harmonic mean of level 3 units across level 4 units (or simple average)
    effect_size: float, default=0.25
        The effect size
    power: float, default=0.8
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    L0: int, default=10
        Our initial guess of our minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0
        Average proportion of level 3 units randomly assigned to treatment within level 4 units
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r22: float, default=0
        Proportion of level 2 variance in the outcome explained by level 2 covariates
    r23: float, default=0
        Proportion of level 3 variance in the outcome explained by level 3 covariates
    r2t4: float, default=0
        Proportion of treatment effect variance among level 4 units explained by level 4 covariates
    g4: int, default=0
        Number of covariates at level 4
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = L0 - g4 - 1
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        L1 = pow(M / effect_size, 2) * (
            rho4 * omega4 * (1 - r2t4)
            + rho3 * (1 - r23) / (p * (1 - p) * K)
            + rho2 * (1 - r22) / (p * (1 - p) * K * J)
            + (1 - rho4 - rho3 - rho2) * (1 - r21) / (p * (1 - p) * K * J * n)
        )
        if abs(L1 - L0) < tol:
            break
        L0 = (L1 + L0) / 2
    sample_size = ceil(L0) if df > 0 else None
    return sample_size

def sample_size_bira2c1(
    n: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.05,
    two_tailed: bool = True,
    J0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    g1: int = 0,
    r21: float = 0,
    print_pretty: bool = True,
) -> Optional[int]:
    """Calculates the minimum sample size of a Two-Level Blocked Individual-level Random Assignment Design,
    Treatment at Level 1

    Parameters
    ----------
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    effect_size: float, default=0.25
        The effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.1
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    J0: int, default=10
        Our initial guess of the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.5
        Average proportion of level 1 units randomly assigned to treatment within level 2 units
    g1: int, default=0
        Number of covariates at level 1
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = J0 * (n - 1) - g1 - 1
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        J1 = pow(M / effect_size, 2) * ((1 - r21) / (p * (1 - p) * n))
        if abs(J1 - J0) < tol:
            break
        J0 = (J1 + J0) / 2
    sample_size = ceil(J0) if df > 0 else None
    if print_pretty:
        print(f"J = {sample_size}")
    return sample_size

def sample_size_bira2f1(
    n: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.05,
    two_tailed: bool = True,
    J0: int = 10,
    tol: bool = 0.10,
    p: float = 0.50,
    g1: int = 0,
    r21: float = 0,
    print_pretty: bool = True,
) -> Optional[int]:
    """Calculates the minimum sample size of a Two-Level Blocked (Fixed) Individual-level Random Assignment Design,
    Treatment at Level 1

    Parameters
    ----------
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    effect_size: float, default=0.25
        The effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    J0: int, default=10
        Our initial guess of the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.5
        Average proportion of level 1 units randomly assigned to treatment within level 2 units
    g1: int, default=0
        Number of covariates at level 1
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = J0 * (n - 2) - g1
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        J1 = pow(M / effect_size, 2) * ((1 - r21) / (p * (1 - p) * n))
        if abs(J1 - J0) < tol:
            break
        J0 = (J1 + J0) / 2
    sample_size = ceil(J0) if df > 0 else None
    print(f"J = {sample_size}")
    return sample_size

def sample_size_bira2r1(
    rho2: float,
    omega2: float,
    n: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.05,
    two_tailed: bool = True,
    J0: int = 10,
    tol: bool = 0.10,
    p: float = 0.50,
    g2: int = 0,
    r21: float = 0,
    r2t2: float = 0,
    print_pretty: bool = True,
) -> Optional[int]:
    """Calculates the minimum sample size of a Two-Level Individual-level Random Assignment Design,
    Treatment at Level 1

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    omega2: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 2 units to the residual
        variance at level 2
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    effect_size: float, default=0.25
        The effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    J0: int, default=10
        Our initial guess of the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.5
        Average proportion of level 1 units randomly assigned to treatment within level 2 units
    g2: int, default=0
        Number of covariates at level 2
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r2t2: float, default=0
        Proportion of treatment effect variance among level 2 units explained by level 2 covariates
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = J0 - g2 - 1
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        J1 = pow(M / effect_size, 2) * (
            rho2 * omega2 * (1 - r2t2) + (1 - rho2) * (1 - r21) / (p * (1 - p) * n)
        )
        if abs(J1 - J0) < tol:
            break
        J0 = (J1 + J0) / 2
    sample_size = ceil(J0) if df > 0 else None
    if print_pretty:
        print(f"J = {sample_size}")
    return sample_size

def sample_size_bira3r1(
    rho2: float,
    rho3: float,
    omega2: float,
    omega3: float,
    n: float,
    J: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed: bool = True,
    K0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    r21: float = 0,
    r2t2: float = 0,
    r2t3: float = 0,
    g3: int = 0,
    print_pretty: bool = True,
) -> Optional[int]:
    """Calculates the minimum sample size of a Three-Level Blocked Individual-level Random Assignment Design,
    Treatment at Level 1

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    rho3: float
        Proportion of variance in the outcome between level 3 units (unconditional ICC3)
    omega2: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 2 units to the residual
        variance at level 2
    omega3: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 3 units to the residual
        variance at level 3
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    J: float
        Harmonic mean of level 2 units across level 3 units (or simple average)
    effect_size: float, default=0.25
        The effect size
    power: float, default=0.8
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    K0: int, default=10
        Our initial guess to the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0
        Average proportion of level 1 units randomly assigned to treatment within level 2 units
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r2t2: float, default=0
        Proportion of treatment effect variance among level 2 units explained by level 2 covariates
    r2t3: float, default=0
        Proportion of treatment effect variance among level 3 units explaiend by level 3 covariates
    g3: int, default=0
        Number of covariates at level 3
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = K0 - g3 - 1
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        K1 = pow(M / effect_size, 2) * (
            rho3 * omega3 * (1 - r2t3)
            + rho2 * omega2 * (1 - r2t2) / J
            + (1 - rho3 - rho2) * (1 - r21) / (p * (1 - p) * J * n)
        )
        if abs(K1 - K0) < tol:
            break
        K0 = (K1 + K0) / 2
    sample_size = ceil(K0) if df > 0 else None
    if print_pretty:
        print(f"K = {sample_size}")
    return sample_size

def sample_size_bira4r1(
    rho2: float,
    rho3: float,
    rho4: float,
    omega2: float,
    omega3: float,
    omega4: float,
    n: float,
    J: float,
    K: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed: bool = True,
    L0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    r21: float = 0,
    r2t2: float = 0,
    r2t3: float = 0,
    r2t4: float = 0,
    g4: int = 0,
    print_pretty: bool = True,
) -> Optional[float]:
    """Calculates the minimum sample size of a Four-Level Blocked Individual-level Random Assignment Design,
    Treatment at Level 1

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    rho3: float
        Proportion of variance in the outcome between level 3 units (unconditional ICC3)
    rho4: float
        Proportion of variance in the outcome between level 4 units (unconditional ICC4)
    omega2: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 2 units to the residual
        variance at level 2
    omega3: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 3 units to the residual
        variance at level 3
    omega4: float
        Treatment effect heterogeneity as ratio of treatment effect variance among level 4 units to the residual
        variance at level 4
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    J: float
        Harmonic mean of level 2 units across level 3 units (or simple average)
    K: float
        Harmonic mean of level 3 units across level 4 units (or simple average)
    effect_size: float, default=0.25
        Effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    L0: int, default=10
        Our initial guess to the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.5
        Average proportion of level 1 units randomly assigned to treatment within level 2 units
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r2t2: float, default=0
        Proportion of treatment effect variance among level 2 units explained by level 2 covariates
    r2t3: float, default=0
        Proportion of treatment effect variance among level 3 units explained by level 3 covariates
    r2t4: float, default=0
        Proportion of treatment effect variance among level 4 units explained by level 4 covariates
    g4: int, default=0
        Number of covariates at level 4
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = L0 - g4 - 1
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        L1 = pow(M / effect_size, 2) * (
            rho4 * omega4 * (1 - r2t4)
            + rho3 * omega3 * (1 - r2t3) / K
            + rho2 * omega2 * (1 - r2t2) / (K * J)
            + (1 - rho4 - rho3 - rho2) * (1 - r21) / (p * (1 - p) * K * J * n)
        )
        if abs(L1 - L0) < tol:
            break
        L0 = (L1 + L0) / 2
    sample_size = ceil(L0) if df > 0 else None
    if print_pretty:
        print(f"L = {sample_size}")
    return sample_size

def sample_size_cra2r2(
    rho2: float,
    n: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed: bool = True,
    J0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    g2: int = 0,
    r21: float = 0,
    r22: float = 0,
    print_pretty: bool = True,
) -> Optional[float]:
    """Calculates the minimum sample size of a Two-level Cluster-randomized Trials to Detect Main, Moderation and
    Mediation Effects

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    effect_size: float, default=0.25
        Effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    J0: int, default=10
        Our initial guess to the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.50
        Proportion of level 2 units randomly assigned to treatment
    g2: int, default=0
        Number of covariates at level 2
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r22: float, default=0
        Proportion of level 2 variance in the outcome explained by level 2 covariates
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = J0 - g2 - 2
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        J1 = pow(M / effect_size, 2) * (
            rho2 * (1 - r22) / (p * (1 - p))
            + (1 - rho2) * (1 - r21) / (p * (1 - p) * n)
        )
        if abs(J1 - J0) < tol:
            break
        J0 = (J1 + J0) / 2
    sample_size = ceil(J0) if df > 0 else None
    if print_pretty:
        print(f"J = {sample_size}")
    return sample_size

def sample_size_cra3r3(
    rho2: float,
    rho3: float,
    n: float,
    J: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed: bool = True,
    K0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    g3: int = 0,
    r21: float = 0,
    r22: float = 0,
    r23: float = 0,
    print_pretty: bool = True,
) -> Optional[float]:
    """Calculates the minimum sample size of a Three-level Cluster-randomized Trials to Detect Main, Moderation,
    and Mediation Effects

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    rho3: float
        Proportion of variance in the outcome between level 3 units (unconditional ICC3)
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    J: float
        Harmonic mean of level 2 units across level 3 units (or simple average)
    effect_size: float, default=0.25
        Effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    K0: int, default=10
        Our initial guess to the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0
        Proportion of level 3 units randomly assigned to treatment
    g3: int, default=0
        Number of covariates at level 3
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r22: float, default=0
        Proportion of level 2 variance in the outcome explained by level 2 covariates
    r23: float, default=0
        Proportion of level 3 variance in the outcome explained by level 3 covariates
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = K0 - g3 - 2
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        K1 = pow(M / effect_size, 2) * (
            rho3 * (1 - r23) / (p * (1 - p))
            + rho2 * (1 - r22) / (p * (1 - p) * J)
            + (1 - rho3 - rho2) * (1 - r21) / (p * (1 - p) * J * n)
        )
        if abs(K1 - K0) < tol:
            break
        K0 = (K1 + K0) / 2
    sample_size = ceil(K0) if df > 0 else None
    if print_pretty:
        print(f"K = {sample_size}")
    return sample_size

def sample_size_cra4r4(
    rho2: float,
    rho3: float,
    rho4: float,
    n: float,
    J: float,
    K: float,
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed: bool = True,
    L0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    r21: float = 0,
    r22: float = 0,
    r23: float = 0,
    r24: float = 0,
    g4: int = 0,
    print_pretty: bool = True,
) -> Optional[float]:
    """Calculates the minimum sample size of a Four-Level Cluster-randomized Trial

    Parameters
    ----------
    rho2: float
        Proportion of variance in the outcome between level 2 units (unconditional ICC2)
    rho3: float
        Proportion of variance in the outcome between level 3 units (unconditional ICC3)
    rho4: float
        Proportion of variance in the outcome between level 4 units (unconditional ICC4)
    n: float
        Harmonic mean of level 1 units across level 2 units (or simple average)
    J: float
        Harmonic mean of level 2 units across level 3 units (or simple average)
    K: float
        Harmonic mean of level 3 units across level 4 units (or simple average)
    effect_size: float, default=0.25
        Effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    L0: int, default=10
        Our initial guess to the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.50
        Proportion of level 4 units randomly assigned to treatment
    r21: float, default=0
        Proportion of level 1 variance in the outcome explained by level 1 covariates
    r22: float, default=0
        Proportion of level 2 variance in the outcome explained by level 2 covariates
    r23: float, default=0
        Proportion of level 3 variance in the outcome explained by level 3 covariates
    r24: float, default=0
        Proportion of level 4 variance in the outcome explained by level 4 covariates
    g4: int, default=0
        Number of covariates at level 4
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = L0 - g4 - 2
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        L1 = pow(M / effect_size, 2) * (
            rho4 * (1 - r24) / (p * (1 - p))
            + rho3 * (1 - r23) / (p * (1 - p) * K)
            + rho2 * (1 - r22) / (p * (1 - p) * K * J)
            + (1 - rho4 - rho3 - rho2) * (1 - r21) / (p * (1 - p) * J * K * n)
        )
        if abs(L1 - L0) < tol:
            break
        L0 = (L1 + L0) / 2
    sample_size = ceil(L0) if df > 0 else None
    if print_pretty:
        print(f"L = {sample_size}")
    return sample_size

def sample_size_ira1r1(
    effect_size: float = 0.25,
    power: float = 0.80,
    alpha: float = 0.10,
    two_tailed=True,
    n0: int = 10,
    tol: float = 0.10,
    p: float = 0.50,
    g1: int = 0,
    r21: float = 0,
    print_pretty: bool = True,
) -> Optional[float]:
    """Calculates the minimum sample size of a Individual-level Random Assignment Design

    Parameters
    ----------
    effect_size: float, default=0.25
        Effect size
    power: float, default=0.80
        The power of our test
    alpha: float, default=0.10
        Probability of Type 1 error
    two_tailed: bool, default=True
        Whether our hypothesis is one tailed or two tailed
    n0: int, default=10
        Our initial guess to the minimum sample size
    tol: float, default=0.10
        Tolerance to end iterative process for finding the minimum sample size
    p: float, default=0.50
        Proportion of units randomly assigned to treatment
    g1: int, default=0
        Number of covariates
    r21: float, default=0
        Proportion of variance in the outcome explained by covariates
    print_pretty: bool, default=True
        Whether we wish to print the results similar to PowerUpR's output

    Returns
    -------
    The minimum sample size of our test
    """
    df = 0
    for i in range(100):
        df = n0 - g1 - 1
        if df < 0 or np.isinf(df):
            break
        T1 = (
            abs(t_dist.ppf(alpha / 2, df)) if two_tailed else abs(t_dist.ppf(alpha, df))
        )
        T2 = abs(t_dist.ppf(power, df))
        M = T1 + T2 if power >= 0.5 else T1 - T2
        n1 = pow(M / effect_size, 2) * ((1 - r21) / (p * (1 - p)))
        if abs(n1 - n0) < tol:
            break
        n0 = (n1 + n0) / 2
    sample_size = ceil(n0) if df > 0 else None
    if print_pretty:
        print(f"n = {sample_size}")
    return sample_size

max_stars_repo_name: ConorMcNamara/pyPowerUp

max_issues_repo_name: ConorMcNamara/pyPowerUp

max_forks_repo_name: ConorMcNamara/pyPowerUp


text:
# 2017 August 06, Tang Li Qun
# Project Standy : Uzu Uzu Index modelling
# 4 layer regression multi layer perceptron (MLP) for modelling the uzu uzu index (UUI) used in project Standy
# feel free to play around with it!

# matrix processing
import pandas as pd
import matplotlib.pyplot as plt
import numpy

# neural network pieces
from keras.models import model_from_json
from keras.models import Sequential
from keras.layers import Dense, Dropout, Embedding, LSTM
from keras.layers import Activation
from keras.layers.normalization import BatchNormalization

# neural network metrics
from sklearn.model_selection import KFold
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error

# misc
import os, pickle

# reference
print(__file__)
directory = os.path.dirname(__file__)

# main
def hackathon_preprocess(dataset):
    """
    extremely shitty, only takes one form of input csv
    :param dataset: dirty dataset
    :return: cleaned dataset
    """
    dataset = dataset.ix[:, 2:]  # takes out first column (previous index number) and second column (datse)
    dataset_a = dataset.ix[:, 0]  # +below: takes out atmospheric pressure
    dataset_b = dataset.ix[:, 2:]
    dataset = pd.concat([dataset_a, dataset_b], axis=1)

    dataset.replace(to_replace=['10-', '0+'], value=0, inplace=True)
    dataset.replace(to_replace=['Mon', 'Tue', 'Wed', 'Thu', 'Fri'], value=0, inplace=True)
    dataset.replace(to_replace=['Sat', 'Sun'], value=1, inplace=True)
    dataset = dataset.fillna(value=0)

    return dataset

def uzu_uzu_index_predictor(train_dataset=None, new_input=False, load_saved_model=True, save_model=True, plot_results=False):
    """
    defaults to loading pretrained model. if load_saved_model is False, it will retrain on the dataset in the accompanying folder.
    either pretrained or new NN will be used to predict the set of new inputs and return an uzu uzu index

    :param train_dataset: FUTURE - feed dataset from outside
    :param new_input: input list for prediction (refer to uui_calculator for details)
    :param load_saved_model: default = True. if False, will retrain model
    :param save_model: default = True. will only save model if load_saved_model is False and save_model is True
    :param plot_results: use matplotlib to plot the results as a graph (for checking accuracy during training)
    :return: uzu uzu index
    """

    if new_input is not False:
        #print('receiving input: ', new_input)
        col_labels = ['ind',
                      'datetime',
                      'weekday',
                      'atm_pressure',
                      'cloud_cover',
                      'dew_point',
                      'rain_mm',
                      'rh',
                      'snow_cover',
                      'sunlight_hrs',
                      'temp_celsius',
                      'visibility',
                      'weather_rating',
                      'forest area coverage',
                      'LowPressure']
        new_input.insert(0, 0)
        new_input.insert(0, 0)
        new_input.insert(3, 0)
        new_input = pd.DataFrame(data=[new_input], columns=col_labels)
        new_input = hackathon_preprocess(new_input)
        # show processed inputs
        print('\nthis is your input :')
        print(new_input)

    # if True, load pretrained model
    if load_saved_model:
        # load json and create model
        json_file = open(directory + r'/trained_nn/model.json', 'r')
        loaded_model_json = json_file.read()
        json_file.close()
        model = model_from_json(loaded_model_json)
        # load weights into new model
        model.load_weights(directory + r"/trained_nn/model.h5")
        scaler_x = pickle.load(open(directory + r'/trained_nn/model_datascaler.p', 'rb'))
        scaler_y = pickle.load(open(directory + r'/trained_nn/model_datascaler_y.p', 'rb'))
        print("Loaded model from disk")
        model.compile(loss='mse', optimizer='adam')

        # prediction inputs required
        print('prediction inputs required:')
        #temp = '278,09/08/2016 15:00,Tue,998,6,19.6,0,38,,2.47,36.2,20,2,0.087407407,0.900249511,5.810333102'.split(',')
        print("""day of the week = 'ddd'
        cloud cover
        dew point
        rain
        relative humidity
        snow cover
        sunlight hours
        temperature
        weather visibility
        weather rating
        forest area
        low pressure index""")

    # else do the full training
    else:

        dataset = pd.read_csv(directory + r'/ml_dataset/tokyo_dataset_uui2.csv')
        dataset = hackathon_preprocess(dataset)

        row_len, col_len = dataset.shape
        feature_len = col_len - 1
        print('number of data points:', row_len)
        print('number of features considered', feature_len)

        # fix random seed for reproducibility
        seed = 7
        numpy.random.seed(seed)

        # get train and test ratio
        # split dataset into sets
        # might need padding if the length is different
        ratio = 0.8  # train/test ratio
        length = len(dataset.index)
        split_i = ratio * length

        x_train = dataset.ix[:split_i, :-1]
        y_train = dataset.ix[:split_i, -1]

        x_test = dataset.ix[split_i+1:, :-1]
        y_test = dataset.ix[split_i+1:, -1]

        # turn dataframe into numpy arrays
        # required for sklearn stuff
        x_train = x_train.as_matrix()
        y_train = y_train.as_matrix()

        x_test = x_test.as_matrix()
        y_test = y_test.as_matrix()

        # normalize data
        # fit scaler
        scaler_x = MinMaxScaler().fit(x_train)
        scaler_y = MinMaxScaler().fit(y_train)
        x_train = scaler_x.transform(x_train)
        x_test = scaler_x.transform(x_test)
        y_train = scaler_y.transform(y_train)
        y_test = scaler_y.transform(y_test)

        # save scaler for transforming new datasets
        pickle.dump(scaler_x, open(directory + r'/trained_nn/model_datascaler.p', 'wb'))
        pickle.dump(scaler_y, open(directory + r'/trained_nn/model_datascaler_y.p', 'wb'))

        # neural networks to use

        def regression_nn():
            # define regression model
            # create model
            model = Sequential()

            # input layer + 20 neuron hidden layer
            model.add(Dense(20, input_dim=feature_len, kernel_initializer='normal'))
            model.add(Activation('relu'))

            # 30 neuron hidden layer
            model.add(Dense(30, kernel_initializer='normal'))
            #model.add(BatchNormalization())
            model.add(Activation('relu'))

            # final nonlinear bit
            model.add(Dense(1, kernel_initializer='normal'))
            model.add(BatchNormalization())
            model.add(Activation('tanh'))
            # compile
            # adam optimizer is super fast compared to sgd
            model.compile(loss='mse', optimizer='adam')
            return model

        def lstm_embedded_nn():
            """
            NOPE.
            :return:
            """

            # initialize
            model = Sequential()

            model.add(Embedding(feature_len, output_dim=64))
            model.add(LSTM(32))
            model.add(Dropout(0.5))
            model.add(Dense(1, activation='tanh'))

            model.compile(loss='mse',
                          optimizer='adam')
            return model

        model = regression_nn()
        # regression model params
        model.fit(x_train, y_train, epochs=500, batch_size=24, verbose=1)
        # simple lstm model params
        #model.fit(x_train, y_train, epochs=300, batch_size=32, verbose=1)

        # make prediction
        predictions = model.predict(x_test)

        # score prediction for regression
        # sklearn accuracy_score only for classification
        test_mse = mean_squared_error(y_test, predictions)
        print('test mse : ', test_mse)

        # score prediction for lstm

        result_pred = scaler_y.inverse_transform(predictions)  # inverse transform the prediction
        result_truth = y_test  # do not scale the ground truth!!!!!

        output_pred = [x[0] for x in predictions.tolist()]
        output_truth = y_test.tolist()

        avg_accuracy = []
        for i, _ in enumerate(output_pred):
            guess = output_pred[i]
            truth = output_truth[i]
            #print('%.2f' % guess, ' > ', '%.2f' % truth)
            num_distance = abs(guess - truth)
            accuracy = (1 - abs(num_distance - 0)) * 100
            avg_accuracy.append(accuracy)

        print('average accuracy %.2f%%' %(sum(avg_accuracy) / float(len(avg_accuracy))))

        if plot_results:
            # show in a graph how accurate
            plt.plot(result_pred, color="blue")
            plt.plot(result_truth, color="green")
            plt.show()

    # save trained model
    # serialize model to JSON
    if not load_saved_model:
        if save_model:
            model_json = model.to_json()
            with open(directory + "/trained_nn/model.json", "w") as json_file:
                json_file.write(model_json)
            # serialize weights to HDF5
            model.save_weights(directory + "/trained_nn/model.h5")
            print("Saved model to disk")

    if new_input is not False:
        # make new prediction
        z_test = scaler_x.transform(new_input)
        predictions = model.predict(z_test)
        predictions = scaler_y.inverse_transform(predictions)  # scaled predictions
    else:
        predictions = [[0]]

    predictions = [x[0] for x in predictions.tolist()]
    return predictions

#uui = uzu_uzu_index_predictor(dataset, load_saved_model=True)

ext: py

max_stars_repo_name: panzerstadt/UzuUzuApi

max_issues_repo_name: panzerstadt/UzuUzuApi

max_forks_repo_name: panzerstadt/UzuUzuApi

reason: import numpy


text:
import os
import torch
import numpy as np
import torch.nn as nn
from torch.autograd import Variable
from src.crowd_count import CrowdCounter
from src.data_loader import ImageDataLoader
from src import utils
import h5py
import scipy.io as io
import PIL.Image as Image
import numpy as np
import os
import glob
from matplotlib import pyplot as plt
from scipy.ndimage.filters import gaussian_filter
import scipy
import torchvision.transforms.functional as F
from matplotlib import cm as CM

import torch.backends.cudnn as cudnn
import torch
import torch.nn as nn
import torch.nn.functional as F
import cv2
from tqdm import tqdm
import math
from torchvision import datasets, transforms
from utils_adv_patch import *
import argparse

def array_transpose(array):
    temp = np.transpose(array,(1,2,0))
    return temp

class FC(nn.Module):
    def __init__(self, in_features, out_features, relu=True):
        super(FC, self).__init__()
        self.fc = nn.Linear(in_features, out_features)
        self.relu = nn.ReLU(inplace=True) if relu else None

    def forward(self, x):
        x = self.fc(x)
        if self.relu is not None:
            x = self.relu(x)
        return x

def save_net(fname, net):
    import h5py
    h5f = h5py.File(fname, mode='w')
    for k, v in net.state_dict().items():
        h5f.create_dataset(k, data=v.cpu().numpy())

def load_net(fname, net):
    import h5py
    h5f = h5py.File(fname, mode='r')
    for k, v in net.state_dict().items():
        param = torch.from_numpy(np.asarray(h5f[k]))
        v.copy_(param)

def np_to_variable(x, is_cuda=True, is_training=False, dtype=torch.FloatTensor):
    if is_training:
        v = Variable(torch.from_numpy(x).type(dtype))
    else:
        v = Variable(torch.from_numpy(x).type(dtype), requires_grad=False, volatile=True)
    if is_cuda:
        v = v.cuda()
    return v

def set_trainable(model, requires_grad):
    for param in model.parameters():
        param.requires_grad = requires_grad

def weights_normal_init(model, dev=0.01):
    if isinstance(model, list):
        for m in model:
            weights_normal_init(m, dev)
    else:
        for m in model.modules():
            if isinstance(m, nn.Conv2d):
                # print torch.sum(m.weight)
                m.weight.data.normal_(0.0, dev)
                if m.bias is not None:
                    m.bias.data.fill_(0.0)
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0.0, dev)

def attack(net, tgt_img_var, patch_var, mask_var, patch_init_var, gt_data_var, target_var, criterion):
    net.eval()
    adv_tgt_img_var = torch.mul((1 - mask_var), tgt_img_var) + torch.mul(mask_var, patch_var)

    # while loss > 0.1 in optical flow attack

    for i in range(args.attack_epoch):

        adv_tgt_img_var = Variable(adv_tgt_img_var.data, requires_grad=True)
        just_the_patch_var = Variable(patch_var.data, requires_grad=True)

        adv_out_var = net(adv_tgt_img_var, gt_data_var)
        loss_data = criterion(adv_out_var, target_var)
        loss_reg = F.l1_loss(torch.mul(mask_var, just_the_patch_var), torch.mul(mask_var, patch_init_var))
        loss = (1 - args.alpha) * loss_data + args.alpha * loss_reg
        loss.backward()

        adv_tgt_img_grad = adv_tgt_img_var.grad.clone()
        adv_tgt_img_var.grad.data.zero_()
        patch_var -= torch.clamp(0.5 * args.lr * adv_tgt_img_grad, -2, 2)

        adv_tgt_img_var = torch.mul((1 - mask_var), tgt_img_var) + torch.mul(mask_var, patch_var)
        adv_tgt_img_var = torch.clamp(adv_tgt_img_var, -1, 1)

    return adv_tgt_img_var, patch_var, adv_out_var

def train(net, patch, patch_shape, mask, patch_init, data_loader_train, criterion):

    patch_shape_orig = patch_shape

    for blob in data_loader_train:
        im_data = blob['data']  # (1,1,645,876)
        gt_data = blob['gt_density']  # (1,1,327,546)
        full_imgname = blob['fname']

        data_shape = im_data.shape  # (1,1,786,1024)
        # beta = np.random.randn(data_shape)
        im_data_gt = torch.from_numpy(im_data)
        tgt_img_var = Variable(im_data_gt.cuda())

        gt_data_var = torch.from_numpy(gt_data)
        gt_data_var = Variable(gt_data_var.cuda())

        if args.patch_type == 'circle':
            patch, mask, patch_init, rx, ry, patch_shape = circle_transform(patch, mask, patch_init, data_shape, patch_shape)
        elif args.patch_type == 'square':
            patch, mask, patch_init, rx, ry = square_transform(patch, mask, patch_init, data_shape, patch_shape)

        patch, mask = torch.FloatTensor(patch), torch.FloatTensor(mask)
        patch_init = torch.FloatTensor(patch_init)

        patch_var, mask_var = Variable(patch.cuda()), Variable(mask.cuda())
        patch_init_var = Variable(patch_init.cuda())

        target_var = Variable(10 * gt_data_var.data.clone(), requires_grad=True).cuda()

        adv_tgt_img_var, patch_var, adv_out_var = attack(net, tgt_img_var, patch_var, mask_var, patch_init_var, gt_data_var, target_var, criterion)
        adv_example = adv_tgt_img_var.data.cpu().numpy()
        # patch_store = patch_var.data.cpu().numpy()

        masked_patch_var = torch.mul(mask_var, patch_var)
        patch = masked_patch_var.data.cpu().numpy()
        mask = mask_var.data.cpu().numpy()
        patch_init = patch_init_var.data.cpu().numpy()
        new_patch = np.zeros(patch_shape)
        new_mask = np.zeros(patch_shape)
        new_patch_init = np.zeros(patch_shape)

        for x in range(new_patch.shape[0]):
            for y in range(new_patch.shape[1]):
                new_patch[x][y] = patch[x][y][ry:ry + patch_shape[-2], rx:rx + patch_shape[-1]]
                new_mask[x][y] = mask[x][y][ry:ry + patch_shape[-2], rx:rx + patch_shape[-1]]
                new_patch_init[x][y] = patch_init[x][y][ry:ry + patch_shape[-2], rx:rx + patch_shape[-1]]

        patch = new_patch
        mask = new_mask
        patch_init = new_patch_init

        patch = zoom(patch, zoom=(1, 1, patch_shape_orig[2] / patch_shape[2], patch_shape_orig[3] / patch_shape[3]),
                     order=1)
        mask = zoom(mask, zoom=(1, 1, patch_shape_orig[2] / patch_shape[2], patch_shape_orig[3] / patch_shape[3]),
                    order=0)
        patch_init = zoom(patch_init,
                          zoom=(1, 1, patch_shape_orig[2] / patch_shape[2], patch_shape_orig[3] / patch_shape[3]),
                          order=1)

    return adv_tgt_img_var, patch, adv_out_var, mask, patch_shape

def test(patch, mask, patch_shape, data_loader_val, net):
    mae_gt = 0.0
    mse_gt = 0.0
    mse_adv = 0.0
    mae_adv = 0.0

    net.eval()

    for blob in data_loader_val:
        im_data = blob['data']  # (1,1,645,876)  # np数组
        gt_data = blob['gt_density']  # (1,1,327,546) np数组
        # full_imgname = blob['fname']

        data_shape = im_data.shape  # (1,1,786,1024)

        im_data_gt = torch.from_numpy(im_data)
        tgt_img_var = Variable(im_data_gt.cuda())

        gt_data_var = torch.from_numpy(gt_data)
        gt_data_var = Variable(gt_data_var.cuda())

        density_map = net(tgt_img_var, gt_data_var)

        if args.patch_type == 'circle':
            patch_full, mask_full, _, _, _, _ = circle_transform_test(patch, mask, patch.copy(),
                                                                                data_shape,
                                                                                patch_shape, True)
        elif args.patch_type == 'square':
            patch_full, mask_full, patch_init, rx, ry = square_transform(patch, mask, patch_init, data_shape,
                                                                         patch_shape)

        patch_full, mask_full = torch.FloatTensor(patch_full), torch.FloatTensor(mask_full)
        patch_full, mask_full = patch_full.cuda(), mask_full.cuda()
        patch_var, mask_var = Variable(patch_full), Variable(mask_full)

        adv_tgt_img_var = torch.mul((1 - mask_var), tgt_img_var) + torch.mul(mask_var, patch_var)
        adv_tgt_img_var = torch.clamp(adv_tgt_img_var, -1, 1)

        adv_out_var = net(adv_tgt_img_var, gt_data_var)

        density_map = density_map.data.detach().cpu().numpy()
        adv_out = adv_out_var.data.detach().cpu().numpy()

        gt_count = np.sum(gt_data)
        et_count = np.sum(density_map)
        adv_count = np.sum(adv_out)

        mae_gt += abs(gt_count - et_count)
        mse_gt += ((gt_count - et_count)*(gt_count - et_count))

        mae_adv += abs(gt_count - adv_count)
        mse_adv += ((gt_count - adv_count)*(gt_count - adv_count))

    mae_gt = mae_gt / data_loader_val.get_num_samples()
    mse_gt = np.sqrt(mse_gt / data_loader_val.get_num_samples())

    mae_adv = mae_adv / data_loader_val.get_num_samples()
    mse_adv = np.sqrt(mse_adv / data_loader_val.get_num_samples())

    print('\nMAE_gt: %0.2f, MSE_gt: %0.2f' % (mae_gt, mse_gt))
    print('\nMAE_adv: %0.2f, MSE_adv: %0.2f' % (mae_adv, mse_adv))

    f = open('./attack_results/adv_results.txt', 'a')
    f.write('adv_mae: %s \n' % str(mae_adv))
    f.write('adv_mse: %s \n' % str(mse_adv))
    f.write('\n')
    f.close()

    f = open('./attack_results/normal_results.txt', 'a')
    f.write('normal_mae: %s \n' % str(mae_gt))
    f.write('normal_mse: %s \n' % str(mse_gt))
    f.write('\n')
    f.close()

reason: import numpy,import scipy,from scipy


text:
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
import sys

"""
    u_embedding: Embedding for center word.
    v_embedding: Embedding for neighbor words.
"""

class RNN(nn.Module):

    def __init__(self, emb_size, emb_dimension, pretrained_emb, output_size,
                 num_layers=2, hidden_size=256, dropout=0.0):
        super(RNN, self).__init__()

        self.hidden_size = hidden_size
        self.sup = num_layers * 2             # num_layers * num_directions

        # Loading pre-trained embeddings
        self.embeddings = nn.Embedding(emb_size, emb_dimension)
        self.embeddings.weight = nn.Parameter(pretrained_emb, requires_grad=False)

        # Bi-directional layer
        self.rnn = nn.RNN(emb_dimension, hidden_size, num_layers=num_layers, bidirectional=True)

        # Dropout layer
        self.dropout = nn.Dropout(dropout)

        # Linear Layers
        self.fc1 = nn.Linear(self.sup*hidden_size, int(hidden_size/2))
        self.fc2 = nn.Linear(int(hidden_size/2), output_size)

        # Cross-Entropy Loss
        self.softmax = nn.CrossEntropyLoss()

    def linear_block(self, x, linear_layer):
        linear_out = linear_layer(x)
        linear_out = F.relu(linear_out)
        linear_out = self.dropout(linear_out)
        return linear_out

    def forward(self, input_sentences, y_batch):

        # Get embeddings of each of the words in the sentences
        sentence_of_emb = self.embeddings(input_sentences)

        # [sentence_length, batch_size, emb_dim]
        rnn_input = sentence_of_emb.permute(1,0,2)
        h_0 = torch.zeros(self.sup, input_sentences.size(0), self.hidden_size).cuda()

        # out : [sentence_length, batch_size, 2 * hidden_size]
        out, h_n = self.rnn(rnn_input, h_0)

        # h_n : [batch_size, 4 * self.hidden_size]
        h_n = h_n.permute(1, 0, 2)
        h_n = h_n.contiguous().view( h_n.size(0), -1 )

        linear_out_1 = self.linear_block(h_n, self.fc1)

        logits = self.fc2(linear_out_1)

        loss = self.softmax(logits, y_batch)

        return loss, logits

class CNN(nn.Module):

    def __init__(self, emb_size, emb_dimension, output_size, pretrained_emb, dropout=0.0,
                       kernel_sz=3, stride=1, padding=(1,0), out_filters=256):

        super(CNN, self).__init__()

        # Loading pre-trained embeddings
        self.embeddings = nn.Embedding(emb_size, emb_dimension)
        self.embeddings.weight = nn.Parameter(pretrained_emb, requires_grad=False)

        # Convolutional Layer
        self.conv1 = nn.Conv2d(1, out_filters, (kernel_sz, emb_dimension), stride=stride, padding=padding)

        # Dropout Layer
        self.dropout = nn.Dropout(p=dropout)

        # Linear Layers
        self.fc1 = nn.Linear(out_filters, int(out_filters/2) )
        self.fc2 = nn.Linear(int(out_filters/2), output_size)

        # Cross-Entropy Loss
        self.softmax = nn.CrossEntropyLoss()

    # If sentences are shorter than the kernel size, it yields out of bounds err.
    def conv_block(self, x, conv_layer):

        # [batch_size, out_filters, H_new, 1]
        conv_out = conv_layer(x)

        # [batch_size, out_filters, H_new]
        conv_out_act = F.relu(conv_out.squeeze(3))

        # Max-overtime pooling: Take the maximum value for each of the feature maps
        # [batch_size, out_filters]
        conv_out = F.max_pool1d(conv_out_act, conv_out_act.size()[2]).squeeze(2)

        return conv_out

    def linear_block(self, x, linear_layer):
        linear_out = linear_layer(x)
        linear_out = F.relu(linear_out)
        linear_out = self.dropout(linear_out)
        return linear_out

    # neg_u should correspond to the index of the negative samples sampled
    def forward(self, input_sentences, y_target):

        # Get embeddings of each of the words in the sentences
        sentence_of_emb = self.embeddings(input_sentences)

        # Reformat to 2D Convolutional shape [batch_size, 1, sentence_length, emb_dim]
        sentence_of_emb = sentence_of_emb.unsqueeze(1)

        # Convolutional Block: [batch_size, out_filters]
        max_out = self.conv_block(sentence_of_emb, self.conv1)

        linear_out_1 = self.linear_block(max_out, self.fc1)

        logits = self.fc2(linear_out_1)

        loss = self.softmax(logits, y_target)

        return loss, logits


text:
import matplotlib.pyplot as plt
import matplotlib.figure
from mpl_toolkits.axes_grid1 import make_axes_locatable
from mamosa.autotrack.ProbabilisticTracking import MAPTracker
from mamosa.utils.plot_utils import set_int_axis_ticks
import numpy as np
from typing import Tuple

def plot_marginal_likelihoods(
    tracker: MAPTracker,
    iline: int,
    horizon_n: int = 0,
    relative=True,
    xticks=None,
    figsize: Tuple[float, float] = (6.4, 4.8),
) -> matplotlib.figure.Figure:
    """Plot marginal likelihoods for horizon depths.

    Args:
        tracker: MAPTracker instance.
        iline: Which iline to plot marginal likelihoods for.
        horizon_n: Horizon number.
        relative: Plot marginals likelihoods relatively to the largest one. I.e. the
            likelihoods are scaled to [0, 1].
        xticks: X-axis ticks.
        figsize: Figure size.

    Returns:
        Matplotlib figure object.

    """
    likelihoods = tracker.get_marginal_likelihoods(horizon_n, iline, relative=relative)
    fig = plt.figure(figsize=figsize)
    set_int_axis_ticks()
    if xticks is not None:
        plt.xticks(xticks)
    if relative is False:
        # TODO: Log colorscale instead
        likelihoods = np.log10(likelihoods)
        likelihoods[likelihoods == 0] = likelihoods.max()
    im = plt.imshow(likelihoods.T, cmap="Greys")
    plt.xlabel("$x$")
    plt.ylabel("Two-way travel time $s$")
    divider = make_axes_locatable(plt.gca())
    cax = divider.append_axes("right", "5%", pad="3%")
    plt.colorbar(im, cax=cax)
    return fig

def plot_posterior_marginals(
    tracker: MAPTracker,
    iline: int,
    horizon_n: int = 0,
    figsize: Tuple[float, float] = (6.4, 4.8),
):
    """Plot posterior marginals for horizon depths.

    Args:
        tracker: MAPTracker instance.
        iline: Which iline to plot marginal likelihoods for.
        horizon_n: Horizon number.
        figsize: Figure size.

    Returns:
        Matplotlib figure object.

    """
    marginals = tracker.get_marginal_posteriors(iline, horizon_n)
    fig = plt.figure(figsize=figsize)
    set_int_axis_ticks()
    im = plt.imshow(marginals.T, cmap="Greys")
    plt.xlabel("$x$")
    plt.ylabel("Two-way travel time $s$")
    divider = make_axes_locatable(plt.gca())
    cax = divider.append_axes("right", "5%", pad="3%")
    plt.colorbar(im, cax=cax)
    return fig

max_stars_repo_name: sindrehenriksen/mamosa

max_issues_repo_name: sindrehenriksen/mamosa

max_forks_repo_name: sindrehenriksen/mamosa


text:
"""
Utilities to preprocess and visualize 2D MRI liver data: Multiecho and IDEAL
Useful for training segmentation and classification models
"""

import itertools
import os
import json
import nibabel as nib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import random
import scipy.ndimage.morphology
import scipy.spatial.distance as distance
import skimage
import sklearn
import sklearn.metrics

from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from scipy.stats import rankdata

from absl import logging

def histogram_equalize(img):
    img_cdf, bin_centers = skimage.exposure.cumulative_distribution(img, nbins=1024)
    return np.interp(img, bin_centers, img_cdf)

def load_image(image_path, equalize=True, return_nifti=False):
    if os.path.exists(image_path):
        image_nii = nib.load(image_path)
        image_data = image_nii.get_fdata()
        image_slice = image_data[:, :, 0, :]
        if equalize:
            image_slice = histogram_equalize(image_slice)
    else:
        logging.warn("{} does not exist".format(image_path))
        image_slice, image_nii = None, None
    if return_nifti:
        return image_slice, image_nii
    else:
        return image_slice

def load_mask(mask_path):
    """
    Load mask nifti file from mask_path. Data should be [H, W, D=1, C]
    Use dilation-fill-erode to remove holes / concavities from annotation artifacts.
    returns a [H, W] shaped mask
    Reduces away the channel dimension with a max.
    """
    if os.path.exists(mask_path):
        mask_nii = nib.load(mask_path)
        mask_data = np.max(mask_nii.get_fdata()[:, :, 0, :], axis=-1)
        mask_data = scipy.ndimage.binary_dilation(mask_data, iterations=2)
        mask_data = scipy.ndimage.morphology.binary_fill_holes(mask_data.astype(np.int)).astype(np.float64)
        mask_data = scipy.ndimage.morphology.binary_erosion(mask_data, iterations=2)
        mask_slice = np.stack([1 - mask_data[:, :, ],
                               mask_data[:, :, ]], axis=2)
    else:
        logging.warn("{} does not exist".format(mask_path))
        mask_slice = None
    return mask_slice

def preprocess_magnitude_phase(magnitude_raw, phase_raw, magnitude_mult,
                               phase_mult, name=None):
    if magnitude_raw is None:
        return None
    if phase_raw is None:
        return None
    magnitude = magnitude_mult * magnitude_raw
    mM = magnitude.max()
    if mM > 10.:
        logging.warn("({}) final magnitude has large max magnitude: {}".format(name, mM))
    if mM < .1:
        logging.warn("({}) final magnitude has small max magnitude: {}".format(name, mM))
    phase = phase_mult * phase_raw
    pM, pm = phase.max(), phase.min()
    if not np.isclose(pM - pm, np.pi * 2, rtol=1e-02, atol=1e-02):
        logging.warn("processed phase has support not close to 2pi: {}, {}".format(pm, pM))

    phase_sin = np.sin(phase)
    phase_cos = np.cos(phase)
    try:
        magnt_psin_pcos = np.concatenate([magnitude, phase_sin, phase_cos], axis=-1)
    except ValueError:
        logging.error("({}) failed to concat magnitude, phase_sin, phase_cos. {} {} {}".format(
            name, repr(magnitude.shape), repr(phase_sin.shape), repr(phase_cos.shape)))
        return None

    return magnt_psin_pcos

def plot_mask_image(image_slice, mask_slice):
    """
    visual diagnostics combining an image with its liver mask
    """
    if (image_slice is not None) and (mask_slice is not None):
        fig, axes = plt.subplots(1, 4)
        plt.axis('off')
        both = np.concatenate([image_slice, mask_slice], axis=2)

        axes[0].imshow(both[:, :, [3, 5, 7]], origin="lower")
        axes[1].imshow(both[:, :, [3, 11, 7]], origin="lower")
        axes[2].imshow(both[:, :, [0, 11, 9]], origin="lower")
        axes[3].imshow(both[:, :, [4, 11, 6]], origin="lower")

    return image_slice, mask_slice

def image_sanity_fail(image, shape, description):
    """
    Sanity check on images: training and testing; shape needs to match.
    description affects the logging, on failure.
    """
    if image is None:
        logging.error("{} : image is None".format(description))
        return True
    elif image.shape != shape:
        logging.error("{} : shape is {}, (expecting {})".format(
            description, repr(image.shape), repr(shape)))
        return True
    else:
        return False

def mask_sanity_fail(mask, shape, description):
    """
    Sanity check on training masks; shape needs to match.
    description affects the logging, on failure.
    """
    if mask is None:
        logging.warn("{} : mask is None".format(description))
        return True
    if mask.shape != shape:
        logging.warn("{} : shape is {}, (expecting {})".format(
            description, repr(mask.shape), repr(shape)))
        return True
    mm = mask[..., 1].mean()
    if mm > .5 or mm < .02:  # mostly should be .07 to .12
        logging.warn("{} : foreground mean {}".format(description, mm))
        return True

def np_rescale(_x, axis):
    M = np.max(_x, axis=axis, keepdims=True)
    m = np.min(_x, axis=axis, keepdims=True)
    d = M - m + 1e-5
    return (_x - m) / d

def rankdata_each(x_hwc):
    h, w, c = x_hwc.shape
    z_hwc = np.stack([np_rescale(rankdata(x_hwc[..., _].flatten()), axis=0)
                      for _ in range(c)], axis=1).reshape(list(x_hwc.shape))
    return z_hwc

def _pca_rank_scale_rgb_tiles(z_img, npc=7, triples=((0, 1, 2), (0, 3, 4), (0, 5, 6), None)):
    """
    z_img is of shape nhwc. nhw dimensions are flattened and the result is fit by a pca.
    triples specified
    """
    pca = PCA(n_components=npc)
    z = z_img.reshape([-1, z_img.shape[-1]])
    z = StandardScaler().fit_transform(z)
    pixel_pca_components = pca.fit_transform(z)
    pixel_pca_df = pd.DataFrame(data=pixel_pca_components, columns=['pc{}'.format(_) for _ in range(npc)])
    pixel_pca_arr = pixel_pca_df.to_numpy()

    tiles = []
    for t in triples:
        if t is None:  # top 3 principal components, without percentile-normalization
            arr = np_rescale(pixel_pca_arr[:, :3], axis=1).reshape(list(z_img.shape)[:-1] + [3, ])
        else:
            assert len(t) == 3
            arr = np.stack([np_rescale(rankdata(pixel_pca_arr[:, _]), axis=0)
                            for _ in t], axis=1).reshape(list(z_img.shape)[:-1] + [3, ])
        tiles.append(arr)
    return tiles

def visualize_head(r, c, vis_dict, image_name, indices=None, titles=None, suptitle=None):
    """
    Creates a figure of r by c subplots. The images are selected from vis_dict.
    The images to plot are keyed by a list of indices.
    titles are keyed in the same nmanner as vis_dict. Save the resulting figure as image_name.
    """
    if indices is None:
        indices = range(min(r * c, len(vis_dict)))
    plt.figure(figsize=(16, 16))

    for j, i in enumerate(indices):
        plt.subplot(r, c, j + 1)
        plt.axis('off')
        if titles is not None:
            plt.title(titles[i])
        plt.imshow(vis_dict[i])
        plt.axis('off')
    if suptitle is not None:
        plt.suptitle(suptitle)
    parent_dir, _ = os.path.split(image_name)
    if not os.path.exists(parent_dir):
        os.makedirs(parent_dir)
    plt.savefig(image_name)
    plt.close()

def get_srnp(gt, pred):
    tn, fp, fn, tp = sklearn.metrics.confusion_matrix(gt, pred).ravel()
    return {'specificity': tn / (tn + fp),
            'recall': tp / (tp + fn),
            'negativepv': tn / (tn + fn),
            'precision': tp / (tp + fp)}

def evaluate_srnpdj(gt, pred, K=25):
    """
    Parameters:
        gt: np.array([N, ...])
            ground truth array of binary labels (integers)
        pred: np.array([, ...])
            same shape as ground truth
    Returns:
        srnpdj_N6: np.array(shape=[N, 6])
            for each paired item in (gt, pred), reteurn 5 evaluation numberse:
            specificity, recall, npv, ppv (prcision), dice-sim, jaccard.
        idx_low_metric_dict: dict(str: [int])
            dictionary of indices with lowest specificities, recall, npv, ppv, and dicesim
    """

    dicesim = []
    jaccard = []
    specificity = []
    recall = []
    negativepv = []
    precision = []

    for _y_true, _y_pred in zip(gt, pred):
        _y_true, _y_pred = _y_true.flatten(), _y_pred.flatten()
        srnp = get_srnp(_y_true, _y_pred)
        specificity.append(srnp['specificity'])
        recall.append(srnp['recall'])
        negativepv.append(srnp['negativepv'])
        precision.append(srnp['precision'])

        dicesim.append(1. - distance.dice(_y_true, _y_pred))
        jaccard.append(sklearn.metrics.jaccard_score(_y_true, _y_pred))

    specificity = np.array(specificity)
    recall = np.array(recall)
    negativepv = np.array(negativepv)
    precision = np.array(precision)
    dicesim = np.array(dicesim)
    jaccard = np.array(jaccard)
    logging.info("dicesim shape: {}".format(repr(dicesim.shape)))

    lowest_specificity = specificity.argsort()[:K]
    lowest_recall = recall.argsort()[:K]
    lowest_npv = negativepv.argsort()[:K]
    lowest_ppv = precision.argsort()[:K]
    lowest_dicesim = dicesim.argsort()[:K]
    lowest_jaccard = jaccard.argsort()[:K]

    srnpdj_N6 = np.stack([specificity, recall, negativepv, precision, dicesim, jaccard], axis=1)

    idx_low_metric_dict = {
        'specificity': lowest_specificity,
        'recall': lowest_recall,
        'npv': lowest_npv,
        'ppv': lowest_ppv,
        'dicesim': lowest_dicesim,
        'jaccard': lowest_jaccard,
    }
    return srnpdj_N6, idx_low_metric_dict

def visualize_predictions(x_img, gt, pred, names, prefix, img_dir, r=None, c=5, rot90=0):
    """
    Creates several figure of r by c subplots. The images are saved into img_dir. rot90 is an
    integer indicating the number of right angle ccw rotations to do. 4 means do nothing.
    prefix is used in the file names of images saved to disk.
    x_img is a bunch of multi channel raw images. We use rankorder and pca to visualize the
    pixels.

    gt and pred are ground truth and prediction masks. We visualize them overlayed on one of the input channels.
    We compute the ids of x_img with the lowest specificity, recall, npv, ppv, and dice.
    We plot the worst cases for each.
    Returns:
        srnpdj_N6: np.array(shape=[N, 6])
            for each paired item in (gt, pred), reteurn 5 evaluation numberse:
            specificity, recall, npv, ppv (prcision), dice-sim, jaccard.
    """

    if r is None:
        r = min(5, int(len(gt) / c))
    K = r * c

    srnpdj_N6, idx_low_metric_dict = evaluate_srnpdj(gt, pred, K=K)

    blue_cidx = 0  # pick any channel
    blue = np_rescale(x_img[:, :, :, blue_cidx:blue_cidx + 1], axis=None)
    red, green = gt, pred
    logging.info("rgb: {} {} {}".format(red.shape, green.shape, blue.shape))
    rgb = np.concatenate([red, green, blue], axis=3)

    vis_idx = set(itertools.chain(*idx_low_metric_dict.values()))
    vis_idx = vis_idx.union(range(r * c))
    visuals_dict = dict()  # compute only the images that score badly on some metric
    for i in vis_idx:
        tiles = _pca_rank_scale_rgb_tiles(x_img[i, :, :, :], triples=((0, 1, 2), (0, 0, 0), None))
        tiles_A = np.concatenate([tiles[0], tiles[1]], axis=0)
        tiles_B = np.concatenate([rgb[i], tiles[2]], axis=0)
        tiles_AB = np.concatenate([tiles_A, tiles_B], axis=1)
        tiles_AB = np.rot90(tiles_AB, k=rot90)
        visuals_dict[i] = tiles_AB

    visualize_head(r, c, visuals_dict, os.path.join(img_dir, 'sseg',
                                                    '{}_sseg.png'.format(prefix)),
                   indices=range(r * c), titles=names, suptitle='{}_{}'.format(prefix, 'samples'))
    for metric_name, low_idx in idx_low_metric_dict.items():
        visualize_head(r, c, visuals_dict, os.path.join(img_dir, 'low_{}'.format(metric_name),
                                                        '{}_low_{}.png'.format(prefix, metric_name)),
                       indices=low_idx, titles=names, suptitle='{}_low_{}'.format(prefix, metric_name))

    low_metric_ids = {k: [names[_] for _ in v] for k, v in idx_low_metric_dict.items()}
    with open(os.path.join(img_dir, '{}_low_metric_ids.json'.format(prefix)), 'w') as json_file:
        json.dump(low_metric_ids, json_file, indent=4)

    plt.figure(figsize=(16, 16))

    for i in range(4 * 5):
        tiles = _pca_rank_scale_rgb_tiles(x_img[i, :, :, :], triples=((0, 1, 2), (0, 3, 4), (0, 0, 0), None))
        tiles_A = np.concatenate([tiles[0], tiles[1]], axis=0)
        tiles_B = np.concatenate([tiles[2], tiles[3]], axis=0)
        tiles_AB = np.concatenate([tiles_A, tiles_B], axis=1)
        tiles_AB = np.rot90(tiles_AB, k=rot90)
        plt.subplot(4, 5, i + 1)
        plt.axis('off')
        plt.title(names[i])
        plt.imshow(tiles_AB[:, :, :])

    plt.savefig(os.path.join(img_dir, '{}_visual.png'.format(prefix)))
    plt.close()
    return srnpdj_N6

def make_thresh_conservative(_x, _t, q=75):  # q = 75: keep 1 quarter of the mask.
    # takes in probabilities _x and float threshold _t. returns _t2, a more stringent threshold.
    _flat = _x.flatten()
    _idx = _flat >= _t
    _t2 = np.percentile(_flat[_idx], q)
    return _t2

# Use for scalar regression models

class ImageMaskNiftiReader(object):
    def __init__(self, magnitude_path_dict, phase_path_dict, mask_path_dict):
        self.magnitude_path_dict = magnitude_path_dict
        self.phase_path_dict = phase_path_dict
        self.mask_path_dict = mask_path_dict

    def get_combined_image_mask(self, subject_id):
        magnitude_path = self.magnitude_path_dict[subject_id]
        phase_path = self.phase_path_dict[subject_id]
        mask_path = self.mask_path_dict[subject_id]

        loaded_magnitude, magnitude_nifti = load_image(magnitude_path, equalize=False, return_nifti=True)
        loaded_phase = load_image(phase_path, equalize=False)
        loaded_magn_sin_cos = preprocess_magnitude_phase(loaded_magnitude, loaded_phase, magnitude_mult=1 / 200.,
                                                         phase_mult=np.pi / 4096., name=subject_id)
        loaded_mask = load_mask(mask_path)
        return combine_image_mask(loaded_magn_sin_cos, loaded_mask, method="concat")

def combine_image_mask(image, mask, method="concat"):
    if method == 'concat':
        return np.concatenate([image, mask], axis=-1)
    elif method == 'fg_only':
        raise NotImplementedError()
    elif method == 'bg_fg':
        bg, fg = np.split(mask, 2, axis=-1)
        bg_im = bg * image
        fg_im = fg * image
        return np.concatenate([bg_im, fg_im], axis=-1)
    else:
        raise ValueError("unknown combine_image_mask method: {}".format(method))

def yield_supervision(ids, im_mask_reader, float_df, isna_df=None, batch_size=4,
                      rand_rot_deg=0, rand_translate=0, skip_all_na=True,
                      shuffle=True, num_loops=None, skip_partial_batch=True,
                      HWC_check=(256, 232, 20), yield_format=None):
    if num_loops is None:
        loop_iter = iter(int, 1)  # infinite loop
    else:
        loop_iter = range(num_loops)

    for _loop in loop_iter:
        im_mask_list, float_list, isna_list, batch_ids = [], [], [], []
        idx = list(range(len(ids)))
        if shuffle:
            random.shuffle(idx)
        for _i in idx:
            curr_id = ids[_i]
            try:
                curr_float = float_df.loc[curr_id].values
                if isna_df is not None:
                    curr_isna = isna_df.loc[curr_id].values
                else:
                    curr_isna = ~np.isfinite(curr_float)
                if np.all(curr_isna) and skip_all_na:
                    # this entry does not have any valid supervision
                    # logging.info('skipping {} -- all supervision na'.format(curr_id))
                    continue
                else:
                    curr_float[curr_isna] = 0.

                im_mask = im_mask_reader.get_combined_image_mask(curr_id)
                if im_mask.shape != HWC_check:
                    logging.warn('invalid im_mask shape for id={} (is {}. expected {})-- skipping item and continuing in generator'.format(
                        curr_id, repr(im_mask.shape), repr(HWC_check)))
                    continue
                # simple data augmentation
                if rand_rot_deg:
                    deg = np.random.uniform(-rand_rot_deg, rand_rot_deg)
                    im_mask = scipy.ndimage.rotate(im_mask, deg, reshape=False)
                if rand_translate:
                    padded = np.pad(im_mask,
                                    ((rand_translate, rand_translate),
                                     (rand_translate, rand_translate), (0, 0)), mode='constant')
                    sh = np.random.randint(0, rand_translate * 2)
                    sw = np.random.randint(0, rand_translate * 2)
                    eh = sh + im_mask.shape[0]
                    ew = sw + im_mask.shape[1]
                    im_mask = padded[sh:eh, sw:ew, :]

            except (FileNotFoundError, ValueError, KeyError) as e:
                logging.warn(repr(e))
                logging.warn('unable to read data for id={} -- skipping item and continuing in generator'.format(curr_id))
                continue
            im_mask_list.append(im_mask)
            float_list.append(curr_float)
            isna_list.append(curr_isna)
            batch_ids.append(curr_id)
            if len(im_mask_list) == batch_size:
                try:
                    batch_im_mask = np.stack(im_mask_list, axis=0)
                    batch_float = np.stack(float_list, axis=0)
                    batch_isna = np.stack(isna_list, axis=0)
                    split_float_targets = [_[:, 0] for _ in
                                           np.split(batch_float, batch_float.shape[1], axis=1)]
                    split_isna_targets = [_[:, 0] for _ in
                                          np.split(1. - batch_isna, batch_float.shape[1], axis=1)]
                except ValueError as e:
                    logging.warn(repr(e))
                    logging.warn('unable to stack or split for ids={} -- skipping batch and continuing in generator'.format(repr(batch_ids)))
                    continue

                inputs = [batch_im_mask, batch_float, batch_isna]
                outputs = split_float_targets
                sample_weights = split_isna_targets
                if yield_format == 'inputs_ids':
                    yield inputs, batch_ids
                else:
                    yield inputs, outputs, sample_weights
                im_mask_list, float_list, isna_list, batch_ids = [], [], [], []
        # after looping through
        if len(im_mask_list) > 0:
            if not skip_partial_batch:
                try:
                    batch_im_mask = np.stack(im_mask_list, axis=0)
                    batch_float = np.stack(float_list, axis=0)
                    batch_isna = np.stack(isna_list, axis=0)
                    split_float_targets = [_[:, 0] for _ in
                                           np.split(batch_float, batch_float.shape[1], axis=1)]
                    split_isna_targets = [_[:, 0] for _ in
                                          np.split(1. - batch_isna, batch_float.shape[1], axis=1)]
                except ValueError as e:
                    logging.warn(repr(e))
                    logging.warn('unable to stack or split for ids={} -- skipping final partial batch and continuing in generator'.format(repr(batch_ids)))
                    continue

                inputs = [batch_im_mask, batch_float, batch_isna]
                outputs = split_float_targets
                sample_weights = split_isna_targets
                if yield_format == 'inputs_ids':
                    yield inputs, batch_ids
                else:
                    yield inputs, outputs, sample_weights
                im_mask_list, float_list, isna_list, batch_ids = [], [], [], []

max_stars_repo_name: calico/ukbb-mri-sseg

max_issues_repo_name: calico/ukbb-mri-sseg

max_forks_repo_name: calico/ukbb-mri-sseg


text:
"""Inception Layer."""

import numpy as np

from athenet.layers import Layer, ConvolutionalLayer, ReLU, MaxPool, \
    Concatenation

class InceptionLayer(Layer):
    def __init__(self, n_filters, input_layer_name=None, name='inception'):
        """Create inception layer.

        Input for the Inception layer is in the format
        (batch size, number of channels, image height, image width).

        :param int list of length 6 n_filters: Number of filters in
                                               convolutional layers.
        """
        super(InceptionLayer, self).__init__(input_layer_name, name)

        layer_list1 = [
            ConvolutionalLayer(filter_shape=(1, 1, n_filters[0]),
                               name=name+'/1x1_conv1'),
            ReLU(),
        ]
        layer_list2 = [
            ConvolutionalLayer(filter_shape=(1, 1, n_filters[1]),
                               name=name+'/1x1_conv2'),
            ReLU(),
            ConvolutionalLayer(filter_shape=(3, 3, n_filters[2]),
                               padding=(1, 1),
                               name=name+'/3x3_conv'),
            ReLU(),
        ]
        layer_list3 = [
            ConvolutionalLayer(filter_shape=(1, 1, n_filters[3]),
                               name=name+'/1x1_conv3'),
            ReLU(),
            ConvolutionalLayer(filter_shape=(5, 5, n_filters[4]),
                               padding=(2, 2),
                               name=name+'/5x5_conv'),
            ReLU(),
        ]
        layer_list4 = [
            MaxPool(poolsize=(3, 3),
                    stride=(1, 1),
                    padding=(1, 1)),
            ConvolutionalLayer(filter_shape=(1, 1, n_filters[5]),
                               name=name+'/1x1_conv4'),
            ReLU(),
        ]
        self.layer_lists = [layer_list1, layer_list2, layer_list3, layer_list4]
        layers = np.concatenate(self.layer_lists)
        self.convolutional_layers = [layer for layer in layers
                                     if isinstance(layer, ConvolutionalLayer)]
        self.bottom_layers = [layer_list[0] for layer_list in self.layer_lists]
        self.top_layers = [layer_list[-1] for layer_list in self.layer_lists]
        self.concat = Concatenation()

    @property
    def output_shape(self):
        return self.concat.output_shape

    @property
    def input_layer(self):
        return self._input_layer

    @input_layer.setter
    def input_layer(self, input_layer):
        self._input_layer = input_layer
        self.input_shape = input_layer.output_shape

        for bottom_layer in self.bottom_layers:
            bottom_layer.input_layer = input_layer

        for layers in self.layer_lists:
            for layer, prev_layer in zip(layers[1:], layers[:-1]):
                layer.input_layer = prev_layer
        self.concat.input_layers = self.top_layers

        self.output = self.concat.output
        self.train_output = self.concat.train_output

    def set_params(self, params):
        for layer, p in zip(self.convolutional_layers, params):
            layer.set_params(p)

    def get_params(self):
        return [layer.get_params() for layer in self.convolutional_layers]

max_stars_repo_name: heurezjusz/Athena

max_issues_repo_name: heurezjusz/Athenet

max_forks_repo_name: heurezjusz/Athena


text:
# Copyright (c) Facebook, Inc. and its affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# Smallest Delta Pruning Method
# By James Michael Gallagher

# Purpose: The purpose of this pruning method is to investigate the impact of pruning weights that
# change the least amount when compared to their initial values in a lottery ticket experiment.
# The amount of pruning can be adjusted with the pruning_fraction hyperparameter the same as in the sparse_global
# pruning method.

# Requirements:

import dataclasses
import numpy as np

from foundations import hparams
import models.base
from pruning import base
from pruning.mask import Mask

@dataclasses.dataclass
class PruningHparams(hparams.PruningHparams):
    pruning_fraction: float = 0.2
    pruning_layers_to_ignore: str = None

    _name = 'Hyperparameters for Smallest Delta Pruning'
    _description = 'Hyperparameters that modify the way pruning occurs.'
    _pruning_fraction = 'The fraction of additional weights to prune from the network.'
    _layers_to_ignore = 'A comma-separated list of additional tensors that should not be pruned.'

class Strategy(base.Strategy):
    @staticmethod
    def get_pruning_hparams() -> type:
        return PruningHparams

    @staticmethod
    def prune(pruning_hparams: PruningHparams, initial_model: models.base.Model, trained_model: models.base.Model, current_mask: Mask = None):

        current_mask = Mask.ones_like(trained_model).numpy() if current_mask is None else current_mask.numpy()

        # Determine the number of weights that need to be pruned.
        number_of_remaining_weights = np.sum([np.sum(v) for v in current_mask.values()])
        number_of_weights_to_prune = np.ceil(
            pruning_hparams.pruning_fraction * number_of_remaining_weights).astype(int)

        # Determine which layers can be pruned. This will be the same for the initial and current models
        prunable_tensors = set(trained_model.prunable_layer_names)
        if pruning_hparams.pruning_layers_to_ignore:
            prunable_tensors -= set(pruning_hparams.pruning_layers_to_ignore.split(','))

        # Get the current model weights.
        weights = {k: v.clone().cpu().detach().numpy()
                   for k, v in trained_model.state_dict().items()
                   if k in prunable_tensors}

        # Get the initial model weights
        initial_weights = {k: v.clone().cpu().detach().numpy()
                   for k, v in initial_model.state_dict().items()
                   if k in prunable_tensors}

        # Create vectors of the weights and get their difference
        weight_vector = np.concatenate([v[current_mask[k] == 1] for k, v in weights.items()])
        initial_weight_vector = np.concatenate([v[current_mask[k] == 1] for k, v in initial_weights.items()])
        final_weight_vector = weight_vector - initial_weight_vector

        threshold = np.sort(np.abs(final_weight_vector))[number_of_weights_to_prune - 1] # Sort is least to greatest

        # Keep the weights that changed the most compared to their value in the first model
        # Note that the weights get zeroed when they don't change enough compared to the threshold
        # This will affect the math in the first component of the np.where code but this is ok since the current_mask[k] value
        # will already have been zeroed from a prior pass
        new_mask = Mask({k: np.where(np.abs(weights[k] - initial_weights[k]) > threshold, current_mask[k], np.zeros_like(v)) for k, v in weights.items()})

        for k in current_mask:
                if k not in new_mask:
                        new_mask[k] = current_mask[k]

        #print("Number of remaining weights: " + str(number_of_remaining_weights))
        #print("Number of weights to prune: " + str(number_of_weights_to_prune))

        return new_mask


text:
import torch
import numpy as np
from tqdm import tqdm

import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation
from IPython import get_ipython

from human_body_prior.body_model.body_model import BodyModel
from human_body_prior.tools.omni_tools import copy2cpu as c2c

class Ax3DPose(object):
    def __init__(
        self, ax, joints, lcolor="#3498db", rcolor="#e74c3c", ccolor="#2fb551"
    ):
        """
        Create a 3d pose visualizer that can be updated with new poses.

        Args
          ax: 3d axis to plot the 3d pose on
          lcolor: String. Colour for the left part of the body
          rcolor: String. Colour for the right part of the body
        """
        matplotlib.rcParams["animation.embed_limit"] = 200
        self.joints = joints
        # Start and endpoints of our representation
        self.I = np.array(
            [
                0,
                0,
                0,
                1,
                2,
                3,
                4,
                5,
                6,
                7,
                8,
                9,
                9,
                9,
                12,
                13,
                14,
                16,
                17,
                18,
                19,
                20,
                22,
                23,
                20,
                25,
                26,
                20,
                28,
                29,
                20,
                31,
                32,
                20,
                34,
                35,
                21,
                37,
                38,
                21,
                40,
                41,
                21,
                43,
                44,
                21,
                46,
                47,
                21,
                49,
                50,
            ]
        )
        self.J = np.arange(1, 52)
        # Left / right indicator
        self.LR = np.ones(52, dtype=int)
        self.LR[
            [
                0,
                3,
                6,
                9,
                12,
                15,
                17,
                19,
                21,
                22,
                23,
                24,
                25,
                26,
                27,
                28,
                29,
                30,
                31,
                32,
                33,
                34,
                35,
            ]
        ] = 0
        self.LR[[2, 5, 8, 11, 14]] = 2
        self.ax = ax

        vals = np.zeros((52, 3))

        # Make connection matrix
        self.plots = []
        for i in np.arange(len(self.I)):
            x = np.array([vals[self.I[i], 0], vals[self.J[i], 0]])
            y = np.array([vals[self.I[i], 1], vals[self.J[i], 1]])
            z = np.array([vals[self.I[i], 2], vals[self.J[i], 2]])
            if self.LR[i] == 0:
                c = lcolor
            elif self.LR[i] == 1:
                c = rcolor
            else:
                c = ccolor

            self.plots.append(
                self.ax.plot(x, y, z, lw=3, c=c if self.LR[i] else rcolor)
            )

        self.ax.set_xlabel("x")
        self.ax.set_ylabel("y")
        self.ax.set_zlabel("z")

    def update(self, frame, lcolor="#3498db", rcolor="#e74c3c", ccolor="#2fb551"):
        """
        Update the plotted 3d pose.

        Args
          channels: 96-dim long np array. The pose to plot.
          lcolor: String. Colour for the left part of the body.
          rcolor: String. Colour for the right part of the body.
        Returns
          Nothing. Simply updates the axis with the new pose.
        """
        vals = self.joints[frame]
        for i in np.arange(len(self.I)):
            x = np.array([vals[self.I[i], 0], vals[self.J[i], 0]])
            y = np.array([vals[self.I[i], 1], vals[self.J[i], 1]])
            z = np.array([vals[self.I[i], 2], vals[self.J[i], 2]])
            self.plots[i][0].set_xdata(x)
            self.plots[i][0].set_ydata(y)
            self.plots[i][0].set_3d_properties(z)
            if self.LR[i] == 0:
                c = lcolor
            elif self.LR[i] == 1:
                c = rcolor
            else:
                c = ccolor
            self.plots[i][0].set_color(c)

        r = 1 if np.max([vals[0, 0], vals[0, 1], vals[0, 2]]) <= 10 else 1000
        xroot, yroot, zroot = vals[0, 0], vals[0, 1], vals[0, 2]
        self.ax.set_xlim3d([-r + xroot, r + xroot])
        self.ax.set_zlim3d([-r + zroot, r + zroot])
        self.ax.set_ylim3d([-r + yroot, r + yroot])

def play_motion(joints, output="HTML"):
    fig = plt.figure()
    ax = plt.gca(projection="3d")
    ob = Ax3DPose(ax, joints)

    if output == "interactive":
        print("interactive mode")
        for j in range(joints.shape[0]):
            ob.update(j)
            plt.show(block=False)
            fig.canvas.draw()
            plt.pause(1.0 / 120)
        plt.close()
    elif output == "HTML":
        anim = FuncAnimation(
            fig,
            ob.update,
            frames=np.arange(0, joints.shape[0]),
            interval=1000 / 120,
            repeat=False,
        )
        # return anim.to_html5_video()
        plt.close()
        return anim.to_jshtml()

max_stars_repo_licenses: CNRI-Python Xnet X11

max_issues_repo_licenses: CNRI-Python Xnet X11

max_forks_repo_licenses: CNRI-Python Xnet X11