text:
###############################################################################
# WaterTAP Copyright (c) 2021, The Regents of the University of California,
# through Lawrence Berkeley National Laboratory, Oak Ridge National
# Laboratory, National Renewable Energy Laboratory, and National Energy
# Technology Laboratory (subject to receipt of any required approvals from
# the U.S. Dept. of Energy). All rights reserved.
#
# Please see the files COPYRIGHT.md and LICENSE.md for full copyright and license
# information, respectively. These files are also available online at the URL
# "https://github.com/watertap-org/watertap/"
#
###############################################################################

"""
    This test is of the core IDAES components that allow for the declaration
    and usage of solids phases in conjunction with aqueous phases. This test
    is primarily being used to probe for issues that may exist in how IDAES
    handles this new system and properties, since it is brand new to the
    framework. Several tests of the same basic case are observed.

    Case 1: [Combine] A and B are aqueous, declare AB a solid,
                    add precipitation reaction (record IDAES debug)

    Several studies of this case are repeated for numerous inlet values to
    assess IDAES's ability to capture both precipitation and dissolution

    Solubility Reaction:
        AB <---> A + B

    Case 2: Repeat Case 1, but try using new 'log_solubility_product'
"""

# Importing testing libraries
import pytest

# Importing the object for units from pyomo
from pyomo.environ import units as pyunits
from pyomo.environ import Var

# Imports from idaes core
from idaes.core import AqueousPhase, SolidPhase, FlowsheetBlock, EnergyBalanceType
from idaes.core.components import Solvent, Solute, Cation, Anion, Component
from idaes.core.phases import PhaseType as PT

# Imports from idaes generic models
import idaes.generic_models.properties.core.pure.Perrys as Perrys
from idaes.generic_models.properties.core.pure.ConstantProperties import Constant
from idaes.generic_models.properties.core.state_definitions import FTPx, FpcTP
from idaes.generic_models.properties.core.eos.ideal import Ideal

# Importing the enum for concentration unit basis used in the 'get_concentration_term' function
from idaes.generic_models.properties.core.generic.generic_reaction import (
    ConcentrationForm,
)

# Import the object/function for heat of reaction
from idaes.generic_models.properties.core.reactions.dh_rxn import constant_dh_rxn

# Import safe log power law equation
from idaes.generic_models.properties.core.reactions.equilibrium_forms import (
    log_power_law_equil,
    power_law_equil,
)

# Import built-in van't Hoff function
from idaes.generic_models.properties.core.reactions.equilibrium_constant import (
    van_t_hoff,
)

from idaes.generic_models.properties.core.reactions.equilibrium_forms import (
    solubility_product,
    log_solubility_product,
    log_power_law_equil,
)
from idaes.generic_models.properties.core.reactions.equilibrium_constant import (
    ConstantKeq,
)

# Import specific pyomo objects
from pyomo.environ import (
    ConcreteModel,
    SolverStatus,
    TerminationCondition,
    value,
    Suffix,
)

from idaes.core.util import scaling as iscale
from idaes.core.util.initialization import fix_state_vars, revert_state_vars

import idaes.logger as idaeslog

# Import pyomo methods to check the system units
from pyomo.util.check_units import assert_units_consistent

# Import idaes methods to check the model during construction
from idaes.core.util import get_solver
from idaes.core.util.model_statistics import (
    degrees_of_freedom,
    fixed_variables_set,
    activated_constraints_set,
    number_variables,
    number_total_constraints,
    number_unused_variables,
)

# Import the idaes objects for Generic Properties and Reactions
from idaes.generic_models.properties.core.generic.generic_property import (
    GenericParameterBlock,
)
from idaes.generic_models.properties.core.generic.generic_reaction import (
    GenericReactionParameterBlock,
)

# Import the idaes object for the EquilibriumReactor unit model
from idaes.generic_models.unit_models.equilibrium_reactor import EquilibriumReactor

# Import log10 function from pyomo
from pyomo.environ import log10

# Import scaling helper functions
from watertap.examples.chemistry.chem_scaling_utils import (
    _set_eps_vals,
    _set_equ_rxn_scaling,
    _set_mat_bal_scaling_FpcTP,
    _set_mat_bal_scaling_FTPx,
    _set_ene_bal_scaling,
)

__author__ = "Austin Ladshaw"

# Case 1 Config
case1_thermo_config = {
    "components": {
        "H2O": {
            "type": Solvent,
            "valid_phase_types": PT.aqueousPhase,
            "dens_mol_liq_comp": Constant,
            "enth_mol_liq_comp": Constant,
            "cp_mol_liq_comp": Constant,
            "entr_mol_liq_comp": Constant,
            # Parameter data is always associated with the methods defined above
            "parameter_data": {
                "mw": (18.0153, pyunits.g / pyunits.mol),
                "dens_mol_liq_comp_coeff": (55.2, pyunits.kmol * pyunits.m**-3),
                "cp_mol_liq_comp_coeff": (75.312, pyunits.J / pyunits.mol / pyunits.K),
                "enth_mol_form_liq_comp_ref": (0, pyunits.kJ / pyunits.mol),
                "entr_mol_form_liq_comp_ref": (0, pyunits.J / pyunits.K / pyunits.mol),
            },
            # End parameter_data
        },
        "A": {
            "type": Solute,
            "valid_phase_types": PT.aqueousPhase,
            "dens_mol_liq_comp": Constant,
            "enth_mol_liq_comp": Constant,
            "cp_mol_liq_comp": Constant,
            "entr_mol_liq_comp": Constant,
            "parameter_data": {
                "mw": (18, pyunits.g / pyunits.mol),
                "dens_mol_liq_comp_coeff": (55.2, pyunits.kmol * pyunits.m**-3),
                "cp_mol_liq_comp_coeff": (75.312, pyunits.J / pyunits.mol / pyunits.K),
                "enth_mol_form_liq_comp_ref": (0, pyunits.kJ / pyunits.mol),
                "entr_mol_form_liq_comp_ref": (0, pyunits.J / pyunits.K / pyunits.mol),
            },
        },
        "B": {
            "type": Solute,
            "valid_phase_types": PT.aqueousPhase,
            "dens_mol_liq_comp": Constant,
            "enth_mol_liq_comp": Constant,
            "cp_mol_liq_comp": Constant,
            "entr_mol_liq_comp": Constant,
            "parameter_data": {
                "mw": (18, pyunits.g / pyunits.mol),
                "dens_mol_liq_comp_coeff": (55.2, pyunits.kmol * pyunits.m**-3),
                "cp_mol_liq_comp_coeff": (75.312, pyunits.J / pyunits.mol / pyunits.K),
                "enth_mol_form_liq_comp_ref": (0, pyunits.kJ / pyunits.mol),
                "entr_mol_form_liq_comp_ref": (0, pyunits.J / pyunits.K / pyunits.mol),
            },
        },
        "AB": {
            "type": Component,
            "valid_phase_types": PT.solidPhase,
            "dens_mol_sol_comp": Constant,
            "enth_mol_sol_comp": Constant,
            "cp_mol_sol_comp": Constant,
            "entr_mol_sol_comp": Constant,
            "parameter_data": {
                "mw": (18, pyunits.g / pyunits.mol),
                "dens_mol_sol_comp_coeff": (55.2, pyunits.kmol * pyunits.m**-3),
                "cp_mol_sol_comp_coeff": (75.312, pyunits.J / pyunits.mol / pyunits.K),
                "enth_mol_form_sol_comp_ref": (0, pyunits.kJ / pyunits.mol),
                "entr_mol_form_sol_comp_ref": (0, pyunits.J / pyunits.K / pyunits.mol),
            },
        },
    },
    # End Component list
    "phases": {
        "Liq": {"type": AqueousPhase, "equation_of_state": Ideal},
        "Sol": {"type": SolidPhase, "equation_of_state": Ideal},
    },
    # Default for testing = FpcTP
    "state_definition": FpcTP,
    # "state_definition": FTPx,
    "state_bounds": {
        # "flow_mol": (0, 50, 100),
        "temperature": (273.15, 300, 650),
        "pressure": (5e4, 1e5, 1e6),
    },
    "pressure_ref": 1e5,
    "temperature_ref": 300,
    "base_units": {
        "time": pyunits.s,
        "length": pyunits.m,
        "mass": pyunits.kg,
        "amount": pyunits.mol,
        "temperature": pyunits.K,
    },
}
# End thermo_config definition

# Actual solubility product
reaction_solubility = {
    "base_units": {
        "time": pyunits.s,
        "length": pyunits.m,
        "mass": pyunits.kg,
        "amount": pyunits.mol,
        "temperature": pyunits.K,
    },
    "equilibrium_reactions": {
        "AB_Ksp": {
            "stoichiometry": {("Sol", "AB"): -1, ("Liq", "A"): 1, ("Liq", "B"): 1},
            "heat_of_reaction": constant_dh_rxn,
            "equilibrium_constant": van_t_hoff,
            "equilibrium_form": solubility_product,
            "concentration_form": ConcentrationForm.moleFraction,
            "parameter_data": {
                "dh_rxn_ref": (0.0, pyunits.J / pyunits.mol),
                "k_eq_ref": (10**-10, pyunits.dimensionless),
                "T_eq_ref": (300.0, pyunits.K),
                "reaction_order": {("Sol", "AB"): 0, ("Liq", "A"): 1, ("Liq", "B"): 1},
            }
            # End parameter_data
        }
        # End Reaction
    }
    # End equilibrium_reactions
}
# End reaction_config definition

# Actual solubility product
reaction_log_solubility = {
    "base_units": {
        "time": pyunits.s,
        "length": pyunits.m,
        "mass": pyunits.kg,
        "amount": pyunits.mol,
        "temperature": pyunits.K,
    },
    "equilibrium_reactions": {
        "AB_Ksp": {
            "stoichiometry": {("Sol", "AB"): -1, ("Liq", "A"): 1, ("Liq", "B"): 1},
            "heat_of_reaction": constant_dh_rxn,
            "equilibrium_constant": van_t_hoff,
            "equilibrium_form": log_solubility_product,
            "concentration_form": ConcentrationForm.moleFraction,
            "parameter_data": {
                "dh_rxn_ref": (0.0, pyunits.J / pyunits.mol),
                "k_eq_ref": (10**-10, pyunits.dimensionless),
                "T_eq_ref": (300.0, pyunits.K),
                "reaction_order": {("Sol", "AB"): 0, ("Liq", "A"): 1, ("Liq", "B"): 1},
            }
            # End parameter_data
        }
        # End Reaction
    }
    # End equilibrium_reactions
}
# End reaction_config definition

# Get default solver for testing
solver = get_solver()

def run_case1(xA, xB, xAB=1e-25, scaling=True, rxn_config=None):
    print("==========================================================================")
    print("Case 1: A and B are aqueous, AB is solid that forms from reaction")
    print("xA = " + str(xA))
    print("xB = " + str(xB))
    print("xAB = " + str(xAB))
    print("scaling = " + str(scaling))
    print("including water = " + str(True))
    print()
    model = ConcreteModel()
    model.fs = FlowsheetBlock(default={"dynamic": False})
    model.fs.thermo_params = GenericParameterBlock(default=case1_thermo_config)

    model.fs.rxn_params = GenericReactionParameterBlock(
        default={"property_package": model.fs.thermo_params, **rxn_config}
    )

    model.fs.unit = EquilibriumReactor(
        default={
            "property_package": model.fs.thermo_params,
            "reaction_package": model.fs.rxn_params,
            "has_rate_reactions": False,
            "has_equilibrium_reactions": True,
            "has_heat_transfer": False,
            "has_heat_of_reaction": False,
            "has_pressure_change": False,
            "energy_balance_type": EnergyBalanceType.none,
        }
    )

    total_flow_mol = 10

    # Set flow_mol_phase_comp
    model.fs.unit.inlet.flow_mol_phase_comp[0, "Liq", "A"].fix(xA * total_flow_mol)
    model.fs.unit.inlet.flow_mol_phase_comp[0, "Liq", "B"].fix(xB * total_flow_mol)
    model.fs.unit.inlet.flow_mol_phase_comp[0, "Sol", "AB"].fix(xAB * total_flow_mol)
    model.fs.unit.inlet.flow_mol_phase_comp[0, "Liq", "H2O"].fix(
        (1 - xA - xB - xAB) * total_flow_mol
    )

    model.fs.unit.inlet.pressure.fix(101325.0)
    model.fs.unit.inlet.temperature.fix(298.0)
    model.fs.unit.outlet.temperature.fix(298.0)

    assert degrees_of_freedom(model) == 0

    assert_units_consistent(model)

    # Scaling
    _set_eps_vals(model.fs.rxn_params, rxn_config)
    # NOTE: We skip reaction scaling because we are NOT using the log_solubility_product form in this test
    # _set_equ_rxn_scaling(model.fs.unit, rxn_config)
    _set_mat_bal_scaling_FpcTP(model.fs.unit)

    iscale.calculate_scaling_factors(model.fs.unit)
    assert isinstance(model.fs.unit.control_volume.scaling_factor, Suffix)
    assert isinstance(
        model.fs.unit.control_volume.properties_out[0.0].scaling_factor, Suffix
    )
    assert isinstance(
        model.fs.unit.control_volume.properties_in[0.0].scaling_factor, Suffix
    )
    # End scaling if statement

    solver.options["max_iter"] = 200
    init_options = {**solver.options}
    init_options["bound_relax_factor"] = 1.0e-02
    model.fs.unit.initialize(optarg=init_options, outlvl=idaeslog.DEBUG)

    assert degrees_of_freedom(model) == 0

    solver.options["bound_relax_factor"] = 1.0e-02
    results = solver.solve(model, tee=True)
    del solver.options["bound_relax_factor"]

    assert results.solver.termination_condition == TerminationCondition.optimal
    assert results.solver.status == SolverStatus.ok

    print("comp\toutlet.tot_molfrac")
    for i in model.fs.unit.control_volume.properties_out[0.0].mole_frac_comp:
        print(
            str(i)
            + "\t"
            + str(
                value(
                    model.fs.unit.control_volume.properties_out[0.0].mole_frac_comp[i]
                )
            )
        )
    print()

    # NOTE: Changed all to mole fraction
    for i in model.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp:
        print(
            str(i)
            + "\t"
            + str(
                value(
                    model.fs.unit.control_volume.properties_out[
                        0.0
                    ].mole_frac_phase_comp[i]
                )
            )
        )

    A = value(
        model.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp[
            "Liq", "A"
        ]
    )
    B = value(
        model.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp[
            "Liq", "B"
        ]
    )
    Ksp = value(model.fs.unit.control_volume.reactions[0.0].k_eq["AB_Ksp"].expr)

    print()
    if Ksp * 1.01 >= A * B:
        print("Constraint is satisfied!")
    else:
        print("Constraint is VIOLATED!")
        print("\tRelative error: " + str(Ksp / A / B) + ">=1")
        assert False
    print("Ksp =\t" + str(Ksp))
    print("A*B =\t" + str(A * B))

    print("==========================================================================")

    return model

def run_case2(xA, xB, xAB=1e-25, scaling=True, rxn_config=None, state="FpcTP"):
    print("==========================================================================")
    print(
        "Case 2 (log form): A and B are aqueous, AB is solid that forms from reaction"
    )
    print("xA = " + str(xA))
    print("xB = " + str(xB))
    print("xAB = " + str(xAB))
    print("scaling = " + str(scaling))
    print("including water = " + str(True))
    print()
    model = ConcreteModel()
    model.fs = FlowsheetBlock(default={"dynamic": False})

    if state == "FpcTP":
        case1_thermo_config["state_definition"] = FpcTP
    elif state == "FTPx":
        case1_thermo_config["state_definition"] = FTPx
        case1_thermo_config["state_bounds"]["flow_mol"] = (0, 50, 100)
    else:
        print("Error! Undefined state...")
        assert False

    model.fs.thermo_params = GenericParameterBlock(default=case1_thermo_config)

    model.fs.rxn_params = GenericReactionParameterBlock(
        default={"property_package": model.fs.thermo_params, **rxn_config}
    )

    model.fs.unit = EquilibriumReactor(
        default={
            "property_package": model.fs.thermo_params,
            "reaction_package": model.fs.rxn_params,
            "has_rate_reactions": False,
            "has_equilibrium_reactions": True,
            "has_heat_transfer": False,
            "has_heat_of_reaction": False,
            "has_pressure_change": False,
            "energy_balance_type": EnergyBalanceType.none,
        }
    )

    total_flow_mol = 10

    # Set flow_mol_phase_comp
    if case1_thermo_config["state_definition"] == FpcTP:

        model.fs.unit.inlet.flow_mol_phase_comp[0, "Liq", "A"].fix(xA * total_flow_mol)
        model.fs.unit.inlet.flow_mol_phase_comp[0, "Liq", "B"].fix(xB * total_flow_mol)
        model.fs.unit.inlet.flow_mol_phase_comp[0, "Sol", "AB"].fix(
            xAB * total_flow_mol
        )
        model.fs.unit.inlet.flow_mol_phase_comp[0, "Liq", "H2O"].fix(
            (1 - xA - xB - xAB) * total_flow_mol
        )

    if case1_thermo_config["state_definition"] == FTPx:
        model.fs.unit.inlet.mole_frac_comp[0, "A"].fix(xA)
        model.fs.unit.inlet.mole_frac_comp[0, "B"].fix(xB)
        model.fs.unit.inlet.mole_frac_comp[0, "AB"].fix(xAB)
        model.fs.unit.inlet.mole_frac_comp[0, "H2O"].fix((1 - xA - xB - xAB))
        model.fs.unit.inlet.flow_mol.fix(total_flow_mol)

    model.fs.unit.inlet.pressure.fix(101325.0)
    model.fs.unit.inlet.temperature.fix(298.0)
    model.fs.unit.outlet.temperature.fix(298.0)

    assert degrees_of_freedom(model) == 0

    assert_units_consistent(model)

    # Scaling
    _set_eps_vals(model.fs.rxn_params, rxn_config)
    _set_equ_rxn_scaling(model.fs.unit, rxn_config)
    if case1_thermo_config["state_definition"] == FpcTP:
        _set_mat_bal_scaling_FpcTP(model.fs.unit)
    if case1_thermo_config["state_definition"] == FTPx:
        _set_mat_bal_scaling_FTPx(model.fs.unit)

    iscale.calculate_scaling_factors(model.fs.unit)
    assert isinstance(model.fs.unit.control_volume.scaling_factor, Suffix)

    # Initialize model

    if case1_thermo_config["state_definition"] == FpcTP:
        state_args = {
            "flow_mol_phase_comp": {
                ("Liq", "H2O"): model.fs.unit.inlet.flow_mol_phase_comp[
                    0, "Liq", "H2O"
                ].value,
                ("Liq", "A"): model.fs.unit.inlet.flow_mol_phase_comp[
                    0, "Liq", "A"
                ].value,
                ("Liq", "B"): model.fs.unit.inlet.flow_mol_phase_comp[
                    0, "Liq", "B"
                ].value,
                ("Sol", "AB"): model.fs.unit.inlet.flow_mol_phase_comp[
                    0, "Sol", "AB"
                ].value,
            },
            "pressure": 101325,
            "temperature": 298,
            "flow_mol": 10,
        }

    if case1_thermo_config["state_definition"] == FTPx:
        state_args = {
            "mole_frac_comp": {
                "H2O": model.fs.unit.inlet.mole_frac_comp[0, "H2O"].value,
                "A": model.fs.unit.inlet.mole_frac_comp[0, "A"].value,
                "B": model.fs.unit.inlet.mole_frac_comp[0, "B"].value,
                "AB": model.fs.unit.inlet.mole_frac_comp[0, "AB"].value,
            },
            "pressure": 101325,
            "temperature": 298,
            "flow_mol": 10,
        }

    flags = fix_state_vars(model.fs.unit.control_volume.properties_out, state_args)
    revert_state_vars(model.fs.unit.control_volume.properties_out, flags)

    model.fs.unit.initialize(optarg=solver.options, outlvl=idaeslog.DEBUG)

    assert degrees_of_freedom(model) == 0

    results = solver.solve(model, tee=True)

    assert results.solver.termination_condition == TerminationCondition.optimal
    assert results.solver.status == SolverStatus.ok

    print("comp\toutlet.tot_molfrac")
    for i in model.fs.unit.control_volume.properties_out[0.0].mole_frac_comp:
        print(
            str(i)
            + "\t"
            + str(
                value(
                    model.fs.unit.control_volume.properties_out[0.0].mole_frac_comp[i]
                )
            )
        )
    print()

    # NOTE: Changed all to mole fraction
    for i in model.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp:
        print(
            str(i)
            + "\t"
            + str(
                value(
                    model.fs.unit.control_volume.properties_out[
                        0.0
                    ].mole_frac_phase_comp[i]
                )
            )
        )

    A = value(
        model.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp[
            "Liq", "A"
        ]
    )
    B = value(
        model.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp[
            "Liq", "B"
        ]
    )
    Ksp = value(model.fs.unit.control_volume.reactions[0.0].k_eq["AB_Ksp"].expr)

    print()
    if Ksp * 1.01 >= A * B:
        print("Constraint is satisfied!")
    else:
        print("Constraint is VIOLATED!")
        print("\tRelative error: " + str(Ksp / A / B) + ">=1")
        assert False
    print("Ksp =\t" + str(Ksp))
    print("A*B =\t" + str(A * B))

    print("==========================================================================")

    return model

## ================================= Case 1 Tests ===============================
@pytest.mark.component
def test_case1_low_conc_no_precipitation():
    model = run_case1(
        xA=1e-9, xB=1e-9, xAB=1e-25, scaling=True, rxn_config=reaction_solubility
    )

@pytest.mark.component
def test_case1_mid_conc_no_precipitation():
    model = run_case1(
        xA=1e-9, xB=1e-2, xAB=1e-25, scaling=True, rxn_config=reaction_solubility
    )

@pytest.mark.component
def test_case1_high_conc_with_precipitation():
    model = run_case1(
        xA=1e-2, xB=1e-2, xAB=1e-25, scaling=True, rxn_config=reaction_solubility
    )

@pytest.mark.component
def test_case1_low_conc_with_dissolution():
    model = run_case1(
        xA=1e-9, xB=1e-9, xAB=1e-2, scaling=True, rxn_config=reaction_solubility
    )

@pytest.mark.component
def test_case1_mid_conc_with_dissolution():
    model = run_case1(
        xA=1e-9, xB=1e-2, xAB=1e-2, scaling=True, rxn_config=reaction_solubility
    )

@pytest.mark.component
def test_case1_high_conc_for_all():
    model = run_case1(
        xA=1e-2, xB=1e-2, xAB=1e-2, scaling=True, rxn_config=reaction_solubility
    )

## ================================= Case 2 Tests ===============================
@pytest.mark.component
def test_case2_low_conc_no_precipitation():
    model = model = run_case2(
        xA=1e-9, xB=1e-9, xAB=1e-25, scaling=True, rxn_config=reaction_log_solubility
    )

@pytest.mark.component
def test_case2_mid_conc_no_precipitation():
    model = model = run_case2(
        xA=1e-9, xB=1e-2, xAB=1e-25, scaling=True, rxn_config=reaction_log_solubility
    )

@pytest.mark.component
def test_case2_high_conc_with_precipitation():
    model = model = run_case2(
        xA=1e-2, xB=1e-2, xAB=1e-25, scaling=True, rxn_config=reaction_log_solubility
    )

@pytest.mark.component
def test_case2_low_conc_with_dissolution():
    model = run_case2(
        xA=1e-9, xB=1e-9, xAB=1e-2, scaling=True, rxn_config=reaction_log_solubility
    )

@pytest.mark.component
def test_case2a_mid_conc_with_dissolution():
    model = model = run_case2(
        xA=1e-9, xB=1e-2, xAB=1e-2, scaling=True, rxn_config=reaction_log_solubility
    )

@pytest.mark.component
def test_case2b_mid_conc_with_dissolution():
    model = model = run_case2(
        xA=1e-2, xB=1e-9, xAB=1e-2, scaling=True, rxn_config=reaction_log_solubility
    )

@pytest.mark.component
def test_case2_high_conc_for_all():
    model = model = run_case2(
        xA=1e-2, xB=1e-2, xAB=1e-2, scaling=True, rxn_config=reaction_log_solubility
    )

max_stars_repo_name: srikanthallu/watertap

max_issues_repo_name: srikanthallu/watertap

max_forks_repo_name: srikanthallu/watertap

reason: from pyomo


text:
#! /usr/bin/python3

from sklearn.datasets import make_blobs
import matplotlib.pyplot as plt
import numpy as np

ext: py

max_stars_repo_name: PeterWrighten/PetersKaggle

max_issues_repo_name: PeterWrighten/PetersKaggle

max_forks_repo_name: PeterWrighten/PetersKaggle

reason: import numpy


text:
import numpy as np
from scipy.stats import multivariate_normal
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import cv2

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from image_utils import rgb2gray, pre_process, linear_mapping, window_func_2d, random_warp

DEBUG = True

class SCCF():
    def __init__(self, center_x, center_y, width, height, n, l):
        # Class variables
        self.g_2D = None # Desired response

        self.A = None # Filter numerator
        self.B = None # Filter denominator

        self.DX = []
        self.DY = []
        self.PX = []
        self.PY = []

        self.center_coord = [center_x, center_y] # Center coordinate for target [x, y]
        self.target_size = [width, height] # Size of target [width, height]
        self.search_size = [1.5 * width, 1.5 * height] # Size of search area [width, height]
        self.l = l
        self.n = n

    def get_target_rectangle(self):
        return (self.center_coord[0] - int(self.target_size[0]/2), self.center_coord[1] - int(self.target_size[1]/2), self.target_size[0], self.target_size[1])

    def get_search_rectangle(self):
        return (self.center_coord[0] - int(self.search_size[0]/2), self.center_coord[1] - int(self.search_size[1]/2), int(self.search_size[0]), int(self.search_size[1]))

    def cv2_target_resize(self, x_img):
        return cv2.resize(x_img, (self.target_size[0], self.target_size[1]))

    def cv2_search_resize(self, x_img):
        return cv2.resize(x_img, (int(self.search_size[0]), int(self.search_size[1])))

    def get_target_shape(self):
        return (self.target_size[1], self.target_size[0])

    def get_search_shape(self):
        return (int(self.search_size[1]), int(self.search_size[0]))

    def set_filter(self, target_2D_norm):
        target_2D_pad = np.pad(target_2D_norm, [(int((self.search_size[1] - self.target_size[1])/2), ), (int((self.search_size[0] - self.target_size[0])/2), )], mode='constant')
        target_2D_pad = self.cv2_search_resize(target_2D_pad)
        g = self.g_2D.reshape(-1,1)
        X = np.fft.fft(target_2D_pad.reshape(-1,1))
        X_conj = np.conjugate(X)
        self.A = np.fft.fft(g) * X_conj
        self.B = X * X_conj + self.l

    def initialize(self, init_img):
        target_2D = init_img[self.center_coord[1]-int(self.target_size[1]/2):self.center_coord[1]+int(self.target_size[1]/2),
                        self.center_coord[0]-int(self.target_size[0]/2):self.center_coord[0]+int(self.target_size[1]/2)]

        search_2D = init_img[self.center_coord[1]-int(self.search_size[1]/2):self.center_coord[1]+int(self.search_size[1]/2),
                        self.center_coord[0]-int(self.search_size[0]/2):self.center_coord[0]+int(self.search_size[1]/2)]
        target_2D = rgb2gray(target_2D)
        search_2D = rgb2gray(search_2D)
        target_2D_norm = pre_process(self.cv2_target_resize(target_2D))
        search_2D_norm = pre_process(self.cv2_search_resize(search_2D))

        # Create Desired Responses, Feature * Filters = Responses
        #Parameters to set
        mu = [self.center_coord[0], self.center_coord[1]]
        covariance = [[self.target_size[0]**2/10, 0], [0, self.target_size[1]**2/10]]

        #Create grid and multivariate normal
        y = np.linspace(self.center_coord[1]-int(self.search_size[1]/2),self.center_coord[1]+int(self.search_size[1]/2),search_2D_norm.shape[0])
        x = np.linspace(self.center_coord[0]-int(self.search_size[0]/2),self.center_coord[0]+int(self.search_size[0]/2),search_2D_norm.shape[1])

        X_mesh, Y_mesh = np.meshgrid(x,y)
        pos = np.empty(X_mesh.shape + (2,))
        pos[:, :, 0] = X_mesh
        pos[:, :, 1] = Y_mesh

        r = multivariate_normal(mu, covariance)
        self.g_2D = linear_mapping(r.pdf(pos))
        if DEBUG:
            max_value = np.max(self.g_2D)
            max_pos = np.where(self.g_2D == max_value)

            plt.imshow(self.g_2D, cmap = plt.get_cmap('gray'))
            plt.title('G 2D')
            print(np.median(max_pos[1]))
            print(np.median(max_pos[0]))
            print(max_pos[0])
            print(max_pos)
            plt.plot(np.median(max_pos[1]), np.median(max_pos[0]), 'bo')
            plt.show()

            #Make a 3D plot
            fig = plt.figure()
            ax = fig.gca(projection='3d')
            ax.plot_surface(X_mesh, Y_mesh, self.g_2D,cmap='viridis',linewidth=0)
            ax.set_xlabel('X axis')
            ax.set_ylabel('Y axis'),
            ax.set_zlabel('Z axis')
            plt.show()
        return target_2D, target_2D_norm

    def test(self, frame):
        #target_2D_pad = np.pad(target_2D_norm, [(int((self.search_size[1] - self.target_size[1])/2), ), (int((self.search_size[0] - self.target_size[0])/2), )], mode='constant')
        #target_2D_pad = self.cv2_search_resize(target_2D_pad)
        search_2D = frame[self.center_coord[1]-int(self.search_size[1]/2):self.center_coord[1]+int(self.search_size[1]/2),
                         self.center_coord[0]-int(self.search_size[0]/2):self.center_coord[0]+int(self.search_size[1]/2)]
        search_2D = rgb2gray(search_2D)
        #search_2D_norm = pre_process(self.cv2_search_resize(search_2D))
        search_2D_norm = pre_process(self.cv2_search_resize(search_2D))
        # Filter
        F = self.A / self.B
        f = np.real(np.fft.ifft(np.conjugate(F)))
        X = np.fft.fft(search_2D_norm.reshape(-1,1))
        R = F * X
        x = np.real(np.fft.ifft(X))

        r = np.real(np.fft.ifft(R)).reshape(self.get_search_shape())
        r = linear_mapping(r)

        max_value = np.max(r)
        max_pos = np.where(r == max_value)

        _, (ax1, ax2, ax3) = plt.subplots(1,3)
        ax1.imshow(x.reshape(self.get_search_shape()), cmap = plt.get_cmap('gray'))
        ax1.set_title("Template image x")

        ax2.imshow(f.reshape(self.get_search_shape()), cmap = plt.get_cmap('gray'))
        ax2.set_title("Filter f")

        ax3.imshow(r, cmap = plt.get_cmap('gray'))
        ax3.plot(max_pos[1], max_pos[0], 'bo')
        ax3.set_title("Convolution result r")
        plt.show()

    # pre train the filter on the first frame...
    def train_filter(self, target_2D, target_2D_norm, num_training, rotate = False):
        # pre-process img..
        # pre-process
        # pad
        target_2D_pad = np.pad(target_2D_norm, [(int((self.search_size[1] - self.target_size[1])/2), ), (int((self.search_size[0] - self.target_size[0])/2), )], mode='constant')
        target_2D_pad = self.cv2_search_resize(target_2D_pad)

        G = np.fft.fft(self.g_2D.reshape(-1,1))
        x = target_2D_pad.reshape(-1,1)
        X = np.fft.fft(x)
        X_conj = np.conjugate(X)
        A = G * X_conj
        B = X * X_conj + self.l
        for _ in range(num_training):
            if rotate:
                t_2D_norm = pre_process(self.cv2_target_resize(random_warp(target_2D)))
            else:
                t_2D_norm = pre_process(self.cv2_target_resize(target_2D))

            t_2D_pad = np.pad(t_2D_norm, [(int((self.search_size[1] - self.target_size[1])/2), ), (int((self.search_size[0] - self.target_size[0])/2), )], mode='constant')
            x = self.cv2_search_resize(t_2D_pad).reshape(-1,1)
            X = np.fft.fft(x)
            X_conj = np.conjugate(X)
            A += X_conj * G
            B += X_conj * X  + self.l

        self.A = A
        self.B = B

    def run(self, frame, run_debug = False):
        # Convert image to grayscale
        gray_frame = rgb2gray(frame)

        # Select template
        y_min = np.max([self.center_coord[1] - int(self.target_size[1] / 2),0])
        x_min = np.max([self.center_coord[0] - int(self.target_size[0] / 2),0])
        target_2D = gray_frame[y_min:self.center_coord[1] + int(self.target_size[1] / 2),
                                 x_min:self.center_coord[0] + int(self.target_size[0] / 2)]
        # Preprocess
        target_2D_norm = pre_process(self.cv2_target_resize(target_2D))

        target_2D_pad = np.pad(target_2D_norm, [(int((self.search_size[1] - self.target_size[1])/2), ), (int((self.search_size[0] - self.target_size[0])/2), )], mode='constant')
        target_2D_pad = self.cv2_search_resize(target_2D_pad)

        # Convert to 1-D vector
        x = target_2D_pad.reshape(-1,1)

        # Convolution between test image and filter F(y) = F(p) * F_1
        F = self.A/self.B
        X = np.fft.fft(x)
        Y = X * F

        # Inverse fourier transform to get result
        y = np.real(np.fft.ifft(Y))
        y_2D = y.reshape(self.get_search_shape())
        y_2D = linear_mapping(y_2D)

        # Get translation
        dx, dy, max_value, max_pos = self._get_translation(y_2D)

        if run_debug:
            print("Max_pos_x: {0}".format(max_pos[1]))
            print("Max_pos_y: {0}".format(max_pos[0]))
            print("center_x: {0}".format(self.center_coord[0]))
            print("center_y: {0}".format(self.center_coord[1]))
            print("center_x_i: {0}".format(self.search_size[0]/2))
            print("center_y_i: {0}".format(self.search_size[1]/2))
            print("dx_i: {0}".format(dx))
            print("dy_i: {0}".format(dy))
            print()
            _, (ax1, ax2, ax3, ax4) = plt.subplots(1,4)
            ax1.imshow(target_2D, cmap = plt.get_cmap('gray'))
            ax1.set_title("Template image x")

            ax2.imshow(target_2D_pad.reshape(self.get_search_shape()), cmap = plt.get_cmap('gray'))
            ax2.set_title("Pre processed template image x")

            f = np.real(np.fft.ifft(np.conjugate(F)))
            ax3.imshow(f.reshape(self.get_search_shape()), cmap = plt.get_cmap('gray'))
            ax3.set_title("Filter f")

            ax4.imshow(y_2D, cmap = plt.get_cmap('gray'))
            # ax4.arrow(target_2D_pad.shape[1]/2, target_2D_pad.shape[0]/2, dy, dx,
            #         head_width=0.05, head_length=1, fc='k', color='red')
            ax4.plot(target_2D_pad.shape[1]/2, target_2D_pad.shape[0]/2, 'ro')
            ax4.plot(max_pos[1], max_pos[0], 'bo')
            ax4.set_title("Convolution result y. Max val: {0}".format(max_value))
            plt.show()

        # Update center points of template with momentum
        beta = 0
        self.center_coord[0] = int(self.center_coord[0] + (1 - beta) * dx)
        self.center_coord[1] = int(self.center_coord[1] + (1 - beta) * dy)

        # Save
        self.DX.append(dx)
        self.DY.append(dy)
        self.PX.append(self.center_coord[0])
        self.PY.append(self.center_coord[1])

        # get updated template z
        y_min = np.max([self.center_coord[1] - int(self.target_size[1] / 2),0])
        x_min = np.max([self.center_coord[0] - int(self.target_size[0] / 2),0])

        z_img = gray_frame[y_min:self.center_coord[1] + int(self.target_size[1] / 2),
                           x_min:self.center_coord[0] + int(self.target_size[0] / 2)]
        # Preprocess
        z_norm = pre_process(self.cv2_target_resize(z_img))

        z_2D_pad = np.pad(z_norm, [(int((self.search_size[1] - self.target_size[1])/2), ), (int((self.search_size[0] - self.target_size[0])/2), )], mode='constant')
        z_2D_pad = self.cv2_search_resize(z_2D_pad)

        # Convert to 1-D vector
        z = z_2D_pad.reshape(-1,1)

        Z = np.fft.fft(z)
        Y = X * F

        # Online update
        self._online_update(Z, Y)
        return

    def _get_translation(self, y_2D):
        #max_y, max_x = np.unravel_index(np.argmax(y_2D), y_2D.shape)  # find the match
        max_value = np.max(y_2D)
        max_pos = np.where(y_2D == max_value)
        dx = int(np.mean(max_pos[1]) - self.search_size[0]/2)
        dy = int(np.mean(max_pos[0]) - self.search_size[1]/2)
        return dx, dy, max_value, max_pos

    def _online_update(self, Z, Y):
        #return
        Z_conj = np.conjugate(Z)
        self.A = (1 - self.n) * self.A + self.n * (Y * Z_conj)
        self.B = (1 - self.n) * self.B + self.n * (Z * Z_conj + self.l)
        return

max_stars_repo_name: ngunnar/object-tracking

max_issues_repo_name: ngunnar/object-tracking

max_forks_repo_name: ngunnar/object-tracking

reason: import numpy,from scipy


text:
import cv2
from scipy import spatial
import numpy as np

def compare_images(image1, image2):
    stats = {
        "mse": _mse(image1, image2),
        "cosine": _cosine(image1, image2),
        'template': _match_template(image1, image2)[0][0],
        'histogram': _hist(image1, image2),
        'correlation': _hist_correl(image1, image2),
        'chi_sqr': _hist_chisqr(image1, image2),
        'intersect': _hist_intersetct(image1, image2)

    }

    return stats

def _mse(image1, image2):
    err = np.sum((image1.astype("float") - image2.astype("float")) ** 2)
    err /= float(image1.shape[0] * image1.shape[1])
    return err

def _cosine(image1, image2):
    im1 = image1.flatten()
    im2 = image2.flatten()
    result = spatial.distance.cosine(im2, im1)
    return result

def _match_template(image1, image2):
    res = cv2.matchTemplate(image1, image2, cv2.TM_SQDIFF_NORMED)
    return 1 - res

def _surf(image1, image2):
    surf = cv2.xfeatures2d.SURF_create()

    keypoints1, descriptor1 = surf.detectAndCompute(image1, None)
    keypoints2, descriptor2 = surf.detectAndCompute(image2, None)

    # Match descriptor vectors using FLANN match
    FLANN_INDEX_KDTREE = 1
    flann_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)

    matcher = cv2.BFMatcher()
    matches = matcher.match(descriptor1[1], descriptor2[1])
    matches = sorted(matches, key=lambda val: val.distance)
    distances = [match.distance for match in matches if match.distance < 0.2]

    im_with_k2 = cv2.drawKeypoints(image1, keypoints1, np.array([]), color=255,
                                   flags=cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)

    im_with_k1 = cv2.drawKeypoints(image2, keypoints2, np.array([]), color=255,
                                   flags=cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)

    cv2.imshow("k1", im_with_k1)
    cv2.imshow("k2", im_with_k2)

    return len(distances)

def _hist_correl(image1, image2):
    hist1 = cv2.calcHist([image1], [0], None, [256], [0, 256])
    hist2 = cv2.calcHist([image2], [0], None, [256], [0, 256])
    res = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)
    return res

def _hist_chisqr(image1, image2):
    hist1 = cv2.calcHist([image1], [0], None, [256], [0, 256])
    hist2 = cv2.calcHist([image2], [0], None, [256], [0, 256])
    res = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CHISQR)
    return res

def _hist_intersetct(image1, image2):
    hist1 = cv2.calcHist([image1], [0], None, [256], [0, 256])
    hist2 = cv2.calcHist([image2], [0], None, [256], [0, 256])
    res = cv2.compareHist(hist1, hist2, cv2.HISTCMP_INTERSECT)
    return res

def _hist(image1, image2):
    hist1 = cv2.calcHist([image1], [0], None, [256], [0, 256])
    hist2 = cv2.calcHist([image2], [0], None, [256], [0, 256])
    res = cv2.compareHist(hist1, hist2, cv2.HISTCMP_BHATTACHARYYA)
    return res


text:
import numpy as np
import pandas as pd
import networkx as nx

def load_mats(data, ROIS):
    print(f"Shape of passed ROI-ROI matrix data is: {np.shape(data['corrs'])}")
    print(f"Number of ROI indices passed is: {np.shape(ROIS['ROIS'])}")

    ROI_labels = []
    for i in range(np.shape((ROIS['ROIS']))[1]):
         ROI_labels.append(ROIS['ROIS'][:,i][0][0])

    mats = []
    #iterate over 3rd dimension(each participant, and grab their corr matrix)
    for i in range(data['corrs'].shape[-1]):
        #load correlations for given participant and label according to ROIs
        patient_conn = pd.DataFrame((data['corrs'][:,:,i]), columns = ROI_labels, index = ROI_labels)

        #mark all correlations of ROI to self (originally NaN) as weight 0.0 (standard in the literature)
        patient_conn.fillna(0, inplace = True)

        #future: change to different correlation statistic? Filter out super low correlations? etc.

        #filter for 264 atlas ROIs
        patient_conn = patient_conn.filter(regex='264', axis = 1)
        patient_conn = patient_conn.filter(regex='264', axis = 0)

        mats.append(patient_conn)

    print(f"All matrices filtered for '264' labeled atlas ROIs only. Shape of each participant matrix is {np.shape(mats[0])}.")

    return mats

#--------DYNAMIC THRESHOLDING-----------------------------------------------------------------
#helper function to take values and determine if they are above a cutoff
def pass_through_cutoff(val, cutoff):
    if val <= cutoff:
        val = 0
    return val

def dynamic_threshold(mat, percentile = 75):
    '''Dynamically threshold Fisher's Z values so that only the specified
    percentile and up of correlation values are retained. This gives all graphs a density of
    1-percentile.

    Expects a matrix with negative correlations already removed. Supported by current rsfMRI GT Literature.'''
    arr = mat.values
    cutoff = np.percentile(arr, 75)
    for col in mat.columns:
        mat[col] = mat[col].apply(pass_through_cutoff, args = [cutoff])
    return mat

def threshold_all_mats(mats, percentile):
    '''Threshold all adjacency matrices using dynamic_threshold function'''
    for i in range(len(mats)):
        mats[i] = dynamic_threshold(mats[i], percentile)
    return mats

#---------------------------------------------------------------------------------------------------

#------------------CLEANING/MAKING PRETTY---------------------------------------------------------------
def remove_atlas(str):
    str = str.replace('BN_Atlas_264_2mm_wkbrois.', '')
    return str

#Clean up ROI names
def clean_roi_names(mats):
    for i in range(len(mats)):
        mats[i] = mats[i].rename(remove_atlas, axis=1)
        mats[i] = mats[i].rename(remove_atlas, axis=0)
    return mats

#---------------------------------------------------------------------------------------------------

#-----------------------GRAPH THEORY FUNCTIONS ------------------------------------------------------------

#insert any graph and get node of interest which should work for all graphs (same ROIs)
def get_nodes_of_interest(graph):
    nodes_of_interest = []
    for node in graph.nodes:
        if 'Amy' in node:
            print(node)
            nodes_of_interest.append(node)
    return nodes_of_interest

def generate_data_for_subject(mats, t1_covariates, nodes_of_interest, subject_number):
    '''Generates Graph Theory Data for a given subject.

    Expects list of weighted adjacency matrices for all subjects,
    covariates matrix/array of same length as # of subjects,
    nodes of interest for certain metrics,
    subject number to calculate for.

    Outputs a dataframe with all graph theory data for this subject'''
    #INIT SUBJECT
    subject = pd.Series(dtype = 'int64')

    #GRAB COVARIATES
    cov = t1_covariates.loc[0,:]

    subject['Subject_Number'] = cov['conn.sub.num']
    subject['mTBI'] = cov['mTBI']
    subject['Male'] = cov['male']

    #subject

    #GENERATE GRAPH FOR THIS SUBJECT
    g = nx.from_pandas_adjacency(mats[subject_number])
    g.name = f'subj_{1}'

    #unhash if you want to display patient graph:
    #edges, weights = zip(*nx.get_edge_attributes(g,'weight').items())
    #pos = nx.spring_layout(p1)
    #nx.draw(p1, pos, node_color='black', edgelist=edges, edge_color=weights, width=10.0, edge_cmap=plt.cm.coolwarm)

    #CLUSTERING COEFFICIENTS
    subject['Avg_Clustering (W)'] = nx.average_clustering(g, weight = 'weight')
    subject['Avg_Clustering (UW)'] = nx.average_clustering(g)

    #INVESTIGATE CHARACTERISTIC PATH LENGTH FOR G
    UW_CPL = nx.average_shortest_path_length(g)

    subject['Characteristic_Path_Length'] = UW_CPL

    #INVESTIGATE DENSITY (Should be the same for all graphs!! ~1-percentile chosen for dynamic thresholding)
    subject['Density'] = nx.density(g)

    #INVESTIGATE TRANSITIVITY
    subject['Transitivity'] = nx.transitivity(g)

    #NODES OF INTEREST-----------------------------------------------------------------------------------------------
    #NODAL DEGREE FOR NODES OF INTEREST
    #Investigate Amygdala Nodes Currently (Nodal Degree)
    degree_view = g.degree(nbunch = nodes_of_interest)
    nodes = (node for (node, val) in degree_view)
    for node in nodes:
        subject[f"Node_Degree--{node}"] = degree_view[node]

    #NODAL STRENGTH FOR NODES OF INTEREST
    degree_view = g.degree(nbunch = nodes_of_interest, weight = 'weight')
    nodes = (node for (node, val) in degree_view)
    for node in nodes:
        subject[f"Node_Strength--{node}"] = degree_view[node]

    #CALCULATE CLOSENESS CENTRALITIES FOR NODES OF INTEREST
    dict_ccs = {}
    for n in nodes_of_interest:
        dict_ccs[n] = nx.closeness_centrality(g, n)
        subject[f"Closeness_Centrality--{n}"] = nx.closeness_centrality(g, n)

    #INVESTIGATE SHORTEST PATHS BETWEEN THESE PAIRS OF NODES
    pair_1 = ['L_dAmy', 'L_MB']
    pair_2 = ['L_dAmy', 'R_MB']

    SP_pair1 = nx.shortest_path_length(g,pair_1[0],pair_1[1])
    SP_pair2 = nx.shortest_path_length(g,pair_2[0],pair_2[1])

    subject[f'Shortest_Path--{pair_1[0]}&{pair_1[1]}'] = SP_pair1
    subject[f'Shortest_Path--{pair_2[0]}&{pair_2[1]}'] = SP_pair2

    #-----------------------------------------------------------------------------------------------------------------

    #CONVERT TO INT DATATYPE WHERE APPROPRIATE AND MAKE INTO DF

    subject = pd.DataFrame(subject).transpose()

    int_data_dict = {'Subject_Number': 'int32',
                     'mTBI': 'int32',
                     'Male': 'int32',
                     'Node_Degree--L_Amyg_mAmyg (medial amygdala)':'int32',
                     'Node_Degree--R_Amyg_mAmyg (medial amygdala)':'int32',
                     'Node_Degree--L_Amyg_lAmyg (lateral amygdala)':'int32',
                     'Node_Degree--R_Amyg_lAmyg (lateral amygdala)':'int32',
                     'Node_Degree--L_dAmy':'int32',
                     'Node_Degree--R_dAmy':'int32',
                     'Node_Degree--L_mAmy':'int32',
                     'Node_Degree--R_mAmy':'int32',
                     'Node_Degree--L_vlAmy':'int32',
                     'Node_Degree--R_vlAmy': 'int32',
                     'Shortest_Path--L_dAmy&L_MB':'int32',
                     'Shortest_Path--L_dAmy&R_MB': 'int32'}

    subject = subject.astype(int_data_dict)

    return subject

#----------------------------------------------------------------------------------------------------------------------------------

reason: import numpy,import networkx


text:
# -----------------------------------------------------------------------
# Name: add_label.py
# Author: Yoshimura Naoya
# Created: 2018.01.18
# Copyright:   (c) yoshimura 2017
#
# Purpose: make inputs frames from 200Hz sequence
# -----------------------------------------------------------------------

import pandas as pd
import numpy  as np
import datetime as dt

import argparse

#
# Params
#

""" make paser
"""
def make_parser():
    """Initialize command line arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('--sensor',  required=True,
                        help='seonsor type {acc, gyro}')
    parser.add_argument('--N',  required=True,
                        help='window size')
    parser.add_argument('--sub-id',  required=True,
                        help='a sub_id of the user(int)')
    parser.add_argument('--file-input',  required=True,
                        help='a path to a input file, e.g.: */n64/{mg4_200Hz_sub01_mod{mod}_{sensor}_{axis}.csv')
    parser.add_argument('--file-output',  required=True,
                        help='a path to a output file, e.g.: */200Hz/{prefix}_200Hz_sub{}_mod{mod}_{sensor_{axis}.csv')
    parser.add_argument('--label-info',  default='None',
                        help='a path to label info')
    return parser

""" Add labels
"""
def label_gesture(df, label_id, cat, label, start_time, session):
    """ Add labels to single ``gesture'' sequence
    """
    idx = df[df["start_time"] == start_time].index
    if len(idx) != 1:
        print("Error: too many idx.", idx)
    idx = idx[0]
    df.loc[idx:(idx+64*2-1), "label"] = label
    df.loc[idx:(idx+64*2-1), "label_id"] = label_id
    df.loc[idx:(idx+64*2-1), "session"] = session
    df.loc[idx:(idx+64*2-1), "cat"] = cat
    df.loc[idx:(idx+64*2-1), "no"] = range(64*2)
    return df

def label_other(df, label_id, cat, label, start_time, end_time, session=-1):
    """ Add label to single ``Object'', ``none'', and ``control'' sequence
    """
    # Index
    idx_start = df[df["start_time"] == start_time].index
    idx_end   = df[df["start_time"] == end_time].index
    if (len(idx_start) != 1) or (len(idx_end) != 1):
        print("\nError: too many idx.", idx_start, idx_end, "\n")
    idx_start, idx_end = idx_start[0], idx_end[0]
    df.loc[idx_start:idx_end, "label"] = label
    df.loc[idx_start:idx_end, "label_id"] = label_id
    df.loc[idx_start:idx_end, "session"] = session
    df.loc[idx_start:idx_end, "cat"] = cat
    df.loc[idx_start:idx_end, "no"] = range(idx_end-idx_start+1)
    return df

def add_label(df, df_label, mod, axis, sub_id, sensor, file_output, f_shape=None, label_info='None', debug=True):
    """ Add label to entire DataFrame
    """
    #
    # Add labels
    df["label"] = "none"
    df["label_id"] = "--"
    df["session"] = -1
    df["no"] = -1
    df["cat"] = "none"
    if label_info != 'None':
        for i in range(len(df_label)):
            cols_label = ["label_id", "cat", "label", "start_mod{}".format(mod),"end_mod{}".format(mod),"session"]
            (label_id, cat, label, start_time, end_time, session) = df_label.loc[i, cols_label]
            if cat == "gesture":
                df = label_gesture(df, label_id, cat, label, start_time, session)
            else:
                df = label_other(df,   label_id, cat, label, start_time, end_time, session)
    #
    # Write results
    cols_idx = ["start_time", "sub_id", "label_id", "session", "mod", "cat", "label", "no", "axis"]
    cols_data = ["x_{}".format(j) for j in range(0, f_shape[0]*2)]
    df = df[cols_idx+cols_data]
    if not debug:
        file_name = file_output.format(mod=mod, sensor=sensor, axis=axis)
        df.to_csv(file_name, index=True)
        print(">> Done: Write {}".format(file_name))
    return df

# ------------------------------------------------------------------------
"""
Main
"""
def main():
    parser = make_parser()
    args = parser.parse_args()
    print("Args:", args)

    # Params
    sensor      = args.sensor
    N           = int(args.N)
    sub_id      = args.sub_id
    file_input  = args.file_input
    file_output = args.file_output
    f_shape = (N, 1) # @100 Hz

    #
    # Load label info
    print(">> label_info={}".format(args.label_info))
    if args.label_info != 'None':
        df_label = pd.read_csv(args.label_info)
        # df_label["start"] = pd.to_datetime(df_label["start"], format="%Y%m%d %H:%M:%S.%f")
        df_label["end"] = pd.to_datetime(df_label["end"], format="%Y%m%d %H:%M:%S.%f")
        df_label["sub_id"] = df_label["sub_id"].astype(str).str.zfill(2)
        print(df_label.head())

    #
    # Main
    #
    if not sub_id:
        user = ct.NAME_LIST[args.sub_id]
        sub_id = user["ID"]
    for mod in range(5):
        for axis in ["x", "y", "z"]:
            # Read CSV
            file_name = file_input.format(mod=mod,sensor=sensor, axis=axis)
            print("File name: ", file_name)
            print(">> sub_id = {}, mod = {}".format(sub_id, mod))
            df = pd.read_csv(file_name, index_col=0)
            # df["start_time"] = pd.to_datetime(df["start_time"], format="%Y%m%d %H:%M:%S.%f")
            df["sub_id"] = df["sub_id"].astype(str).str.zfill(2)
            # Add labels
            if args.label_info != 'None':
                df_label_tmp = df_label[df_label["sub_id"] == sub_id]
                df_label_tmp = df_label_tmp.reset_index(drop=True)
            else:
                df_label_tmp = None
            add_label(df, df_label_tmp, mod, axis, sub_id, sensor, file_output, f_shape=f_shape, label_info=args.label_info, debug=False)
            print(">> Success")

# ------------------------------------------------------------------------
if __name__=='__main__':
    main()


text:
from keckdrpframework.primitives.base_primitive import BasePrimitive
from kcwidrp.primitives.kcwi_file_primitives import kcwi_fits_writer

import numpy as np
import pkg_resources
import os
import pandas as pd

class CorrectDefects(BasePrimitive):
    """Remove known bad columns"""

    def __init__(self, action, context):
        BasePrimitive.__init__(self, action, context)
        self.logger = context.pipeline_logger

    def _perform(self):
        self.logger.info("Correcting detector defects")

        # Header keyword to update
        key = 'BPCLEAN'
        keycom = 'cleaned bad pixels?'

        # Create flags for bad columns fixed
        flags = np.zeros(self.action.args.ccddata.data.shape, dtype=np.uint8)

        # Nod and Shuffle?
        if self.action.args.nasmask and self.action.args.numopen > 1:
            nastr = "_nas"
        else:
            nastr = ""

        # Does the defect file exist?
        path = "data/defect_%s_%dx%d%s.dat" % (self.action.args.ampmode.strip(),
                                               self.action.args.xbinsize,
                                               self.action.args.ybinsize, nastr)
        package = __name__.split('.')[0]
        full_path = pkg_resources.resource_filename(package, path)
        number_of_bad_pixels = 0   # count of defective pixels cleaned
        if os.path.exists(full_path):
            self.logger.info("Reading defect list in: %s" % full_path)
            defect_table = pd.read_csv(full_path, sep=r'\s+')
            # range of pixels for calculating good value
            pixel_range_for_good_value = 5
            for index, row in defect_table.iterrows():
                # Get coords and adjust for python zero bias
                x0 = row['X0'] - 1
                x1 = row['X1']
                y0 = row['Y0'] - 1
                y1 = row['Y1']
                # Loop over y range
                for by in range(y0, y1):
                    # sample on low side of bad area
                    values = list(self.action.args.ccddata.data[by,
                                  x0-pixel_range_for_good_value:x0])
                    # sample on high side
                    values.extend(self.action.args.ccddata.data[by,
                                  x1+1:x1+pixel_range_for_good_value+1])
                    # get replacement value
                    good_values = np.nanmedian(np.asarray(values))
                    # Replace baddies with gval
                    for bx in range(x0, x1):
                        self.action.args.ccddata.data[by, bx] = good_values
                        flags[by, bx] += 2
                        number_of_bad_pixels += 1
            self.action.args.ccddata.header[key] = (True, keycom)
            self.action.args.ccddata.header['BPFILE'] = (path, 'defect list')
        else:
            self.logger.error("Defect list not found for %s" % full_path)
            self.action.args.ccddata.header[key] = (False, keycom)

        self.logger.info("Cleaned %d bad pixels" % number_of_bad_pixels)
        self.action.args.ccddata.header['NBPCLEAN'] = \
            (number_of_bad_pixels, 'number of bad pixels cleaned')

        log_string = CorrectDefects.__module__
        self.action.args.ccddata.header['HISTORY'] = log_string
        self.logger.info(log_string)

        # add flags array
        self.action.args.ccddata.mask = flags
        self.action.args.ccddata.flags = flags

        if self.config.instrument.saveintims:
            kcwi_fits_writer(self.action.args.ccddata,
                             table=self.action.args.table,
                             output_file=self.action.args.name,
                             output_dir=self.config.instrument.output_directory,
                             suffix="def")

        return self.action.args
    # END: class CorrectDefects()


text:
# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import matplotlib.pyplot as plt

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

# You can write up to 5GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All"
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
data = pd.read_csv('/kaggle/input/ghouls-goblins-and-ghosts-boo/train.csv.zip', index_col="id")
data
data.info()
data.describe()
data["color"].hist()
data["type"].astype('category')
data["type"] = data["type"].astype('category')
data["color"] = data["color"].astype('category')
data["color"]
for color in data["color"].cat.categories:
    print("Color: ", color.upper())
    for monster_type in data["type"].cat.categories:
        print("\t", monster_type, data.query('color == @color & type == @monster_type')["type"].count())
data.loc[(data["color"] == "green") & (data["type"] == "Ghost"), "type"].count() / data.loc[data["color"] == "green", "type"].count()
#green ghosts.. strange, isn't it?
#color splits monsters into 3 nearly equal parts so I think it is useless
data.drop("color", axis=1, inplace=True)
data
from sklearn.linear_model import SGDClassifier
from sklearn.model_selection import cross_val_score
from sklearn.ensemble import RandomForestClassifier

print(cross_val_score(SGDClassifier(), data.drop('type', axis=1), data["type"], cv=3))
print(cross_val_score(RandomForestClassifier(), data.drop(['type'], axis=1), data["type"], cv=3))
linear = SGDClassifier(early_stopping=True, max_iter=10000, learning_rate='adaptive', eta0=2)
linear.fit(data.drop('type', axis=1), data['type'])
linear.score(data.drop('type', axis=1), data['type'])
test_data = pd.read_csv('/kaggle/input/ghouls-goblins-and-ghosts-boo/test.csv.zip', index_col="id")
test_data
predictions = linear.predict(test_data.drop("color", axis=1))

pd.read_csv('/kaggle/input/ghouls-goblins-and-ghosts-boo/sample_submission.csv.zip')
result = pd.DataFrame({'id': test_data.index,
                       'type': predictions})
result.to_csv("monsters_submission.csv", index=False)
result["type"]

max_stars_repo_name: josepablocam/janus-public

max_issues_repo_name: josepablocam/janus-public

max_forks_repo_name: josepablocam/janus-public


text:
from __future__ import print_function
import sys
from collections import OrderedDict
import numpy as np

def parse_cif(infile):
    if isinstance(infile, file):
        pass
    elif isinstance(infile, str):
        infile = open(infile)
    else:
        raise TypeError("Input should either be a file or string.")

    atom_site = OrderedDict()
    with infile as f:
        for line in f:

            if line.startswith('_atom_site.'):
                words = line.split('.')
                atom_site[words[1].strip()] = []

            if line.startswith('ATOM'):
                words = line.split()
                for key, word in zip(atom_site, words):
                    atom_site[key].append(word)

    natoms = len(atom_site['id'])
    dtype = [('atom_id', np.int64), ('name', np.str_, 4),
             ('resn', np.str_, 4), ('chain', np.str_, 2),
             ('resi', np.int64), ('x', np.float64),
             ('y', np.float64), ('z', np.float64),
             ('occupancy', np.float64), ('bfactor', np.float64),
             ('element', np.str_, 2), ('charge', np.str_, 2),
             ('model', np.int64),
             ]

    cifdata = np.zeros(natoms, dtype=dtype)
    cifdata['atom_id'] = np.asarray(atom_site['id'], dtype=np.int64)
    cifdata['name'] = atom_site['label_atom_id']
    cifdata['resn'] = atom_site['label_comp_id']
    cifdata['chain'] = atom_site['label_asym_id']
    cifdata['resi'] = atom_site['label_seq_id']
    cifdata['x'] = atom_site['Cartn_x']
    cifdata['y'] = atom_site['Cartn_y']
    cifdata['z'] = atom_site['Cartn_z']
    cifdata['occupancy'] = atom_site['occupancy']
    cifdata['bfactor'] = atom_site['B_iso_or_equiv']
    cifdata['element'] = atom_site['type_symbol'].title()
    cifdata['charge'] = atom_site['pdbx_formal_charge']
    cifdata['model'] = atom_site['pdbx_PDB_model_num']

    return cifdata

if __name__=='__main__':
    import sys
    infile = sys.argv[1]
    data = parse_cif(infile)

max_stars_repo_name: haddocking/disvis

max_issues_repo_name: haddocking/disvis

max_forks_repo_name: haddocking/disvis


text:
import json
import pickle
import csv
import sys
import jsonlines
import datetime
import numpy as np
import random
import os
import torch
import pathlib

def seed_everything(seed):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

def create_folders_path(path_name):
    path = pathlib.Path(path_name)
    path.parent.mkdir(parents = True,exist_ok = True)

def loadPickleFile(filepath):
	#print("Loading the pickle file from",filepath,"...........")
	pickle_in = open(filepath,"rb")
	example_dict = pickle.load(pickle_in)
	#print("Loaded the pickle File")
	return example_dict

def openCSVfile(filepath, delimiter = ","):
    """
    Returns the lists for csv file
    """
    with open(filepath,"r") as csvfile:
        rows =  csv.reader(csvfile,delimiter = delimiter)
        return list(rows)

def dumpPickleFile(data,filepath):
	pickle_out = open(filepath,"wb")
	#print("Dumping the Pickle file into ",filepath,"...........")
	pickle.dump(data, pickle_out)
	#print("Dumped the pickle File")
	pickle_out.close()

def dumpJsonFile(dictionary,filepath):
	#print("Dumping a dictionary to filepath",filepath,"...............")
	with open(filepath,"w+") as jsonFile:
		json.dump(dictionary,jsonFile,indent=4,sort_keys =True)
	#print("Dumped Successfully",filepath)

def loadJsonFile(filepath):
	#print("Loading a dictionary to filepath",filepath,"...............")
	dictionary = {}
	with open(filepath) as jsonFile:
		dictionary = json.load(jsonFile)
	#print("Loaded Successfully")
	return dictionary

# Convert timestamp
def convert_timestamp(timestamp):
    date = datetime.datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S')
    date_split = date.split("-");month = int(date_split[1]);year = int(date_split[0])
    return month,year,date

def read_json_objects(filepath):
    objects = []
    with jsonlines.open(filepath) as reader:
        for obj in reader:
            created_timestamp = obj['created_utc']#;updated_timestamp = obj['updated_utc']
            month,year,date = convert_timestamp(created_timestamp)
            objects.append([obj,month,year,date])
    return objects

def get_arguments():
    n = len(sys.argv)
    if n != 4:
        print('Incomplete not executing')
        sys.exit(0)
    print("Total arguments passed:", n)
    print("\nName of Python script:", sys.argv[0])
    print("\nArguments passed:", end = " ")

    model_name = sys.argv[1]
    class_name = sys.argv[2]
    json_filename = sys.argv[3]
    save_filename = sys.argv[4]

    return model_name, json_filename, save_filename


text:
import numpy as np
import matplotlib.pyplot as plt
from diskvert import *

clr = ['#93ABAA', '#09D6AB', '#CB2E13']

dd,pd = col2python('diskd.dat')
dw,pw = col2python('diskw.dat')
dc,pc = col2python('diskc.dat')

kabs = lambda rho, T: pd.kabp0 * rho * T**(-3.5)
yyc = pd.kappa_es * (4 * cgs_boltz * dd['trad']**4) / (cgs_mel * cgs_c**2)
tdumm1 = dd['trad'] + dd['heat'] / (4 * cgs_stef * dd['rho'] * yyc)
yyb = kabs(dd['rho'], tdumm1) * (dd['trad'] + tdumm1) * (dd['trad']**2 + tdumm1**2)
tdumm = dd['trad'] + dd['heat'] / (4 * cgs_stef * dd['rho'] * (yyb + yyc))
pdumm = dd['pgas'] * tdumm / dd['temp']

h0   = dc['h'][::2]
rho  = dc['rho'][::2]
heat = dc['heat'][::2]
trad = dc['trad'][::2]
temp = dc['temp'][::2]
T0 = np.logspace(np.log10(0.31 * pd.teff), np.log10(1e3 * pd.teff), 90)
hg, Tg = np.meshgrid(h0, T0)
bil = np.ndarray((len(T0),len(h0)))

for i in range(len(T0)):
    rho1 = rho
    # rho1 = rho * temp / T0[i]
    yyb = kabs(rho1, T0[i]) * (trad + T0[i]) * (trad**2 + T0[i]**2)
    yyc = pd.kappa_es * (4 * cgs_boltz * trad**4) / (cgs_mel * cgs_c**2)
    cool = 4 * cgs_stef * rho1 * (yyb + yyc) * (T0[i] - trad)
    bil[i,:] = 1 - cool / heat

plt.figure()
plt.pcolor(h0, T0, bil, cmap = 'bil1', vmin = -0.2, vmax = 0.2)
plt.loglog()
plt.xlim(2,2e2)

def plotthis(ax, col):
    ax.set_title(col)
    ax.plot(dd['h'], dd[col], color = clr[0], label = 'DYF')
    ax.plot(dw['h'], dw[col], color = clr[1], label = 'CMP')
    ax.plot(dc['h'], dc[col], color = clr[2], label = 'BIL')
    ax.set_xlim(2,2e2)
    ax.loglog()
    ax.legend()

fig, axes = plt.subplots(2, 2, figsize = (9,8), sharex = True)

plotthis(axes[0,0], 'rho')
plotthis(axes[0,1], 'pgas')
axes[0,1].plot(dd['h'], pdumm, '--', color = '#F450E3')
plotthis(axes[1,0], 'temp')
axes[1,0].plot(dd['h'], tdumm, '--', color = '#F450E3')
plotthis(axes[1,1], 'prad')

plt.show()

max_stars_repo_name: gronki/pydiskvert

max_issues_repo_name: gronki/pydiskvert

max_forks_repo_name: gronki/pydiskvert


text:
import glob
import os
import pandas as pd
import json
import matplotlib.pyplot as plt
from progressbar import progressbar
import numpy as np
from pathlib import Path
import pickle
import logging
from matplotlib import scale
import itertools

plt.set_cmap('jet')
def getLogging():
    logging.basicConfig(
        format='%(asctime)s %(levelname)-8s %(message)s',
        level=logging.INFO,
        datefmt='%Y-%m-%d %H:%M:%S')
    return logging
def read_json(file_path):
    with open(file_path,"r") as f:
        result=json.load(f)
    return result
def find_max_len(result):
    length=1
    for i in result.keys():
        length=max(length,len(result[i]))
    return length
def append_single(result,max_len):
    for i in result.keys():
        if len(result[i])<max_len:
            result[i]=result[i]+[result[i][-1]]*(max_len-len(result[i]))
def write_back(result,file_path):
    with open(file_path,"w") as f:
        json.dump(result,f)
def get_path_sets(root_path,same_length=False,suffix="jjson"):
    paths=glob.glob(root_path+'/**/*.'+suffix, recursive=True)
    # print(paths)
    path_sets=set()
    for path in paths:
        pp=Path(paths[0])
        # path_root=str(pp.parent.parent)
        path_root=os.path.join(*path.split("/")[:-2])
#         print(path_root,path_root,"path_root")
        if same_length:
            result_cur=read_json(path)
            max_len=find_max_len(result_cur)
            append_single(result_cur,max_len)
            write_back(result_cur,path)
        # print(result_cur)
        # print(path_root)
        path_sets.add(path_root)
    return path_sets
def merge_single_experiment_results(root_path,suffix="jjson"):
    paths=glob.glob(root_path+'/**/*.'+suffix, recursive=True)
    # print(paths,root_path)
    df_result_cur=pd.DataFrame()
    for path in paths:
        print(path)
        with open(path, "r") as read_file:
        #     print("Converting JSON encoded data into Python dictionary")
            developer = json.load(read_file)
            # print(path,"path")
            pd_frame=pd.DataFrame(developer).fillna(np.nan)
        df_result_cur=pd.concat([df_result_cur,pd_frame])
#     print(df_result_cur)
    return df_result_cur

def merge_multiple_experiment_results(root_path,suffix="jjson"):
    path_sets=get_path_sets(root_path)
    for path_set in progressbar(path_sets):
        df_result_cur=merge_single_experiment_results(path_set,suffix=suffix)
        df_result_cur.to_csv(path_set+"/result.ccsv")

class AutoVivification(dict):
    """Implementation of perl's autovivification feature."""
    def __getitem__(self, item):
        try:
            return dict.__getitem__(self, item)
        except KeyError:
            value = self[item] = type(self)()
            return value
def get_node(root_path,path):
    # print(root_path,path,"root_path,path")
    all_node=path.split("/")
    start_folder=root_path.split("/")[-1]
    # print(root_path.split("/"),all_node)
    ind=all_node.index(start_folder)
    start_node=all_node[ind+1:]
    # print(start_node)
    return start_node
def set_node_val(node_list,multi_level_dict,val):
    for i in node_list[:-1]:
        multi_level_dict=multi_level_dict[i]
    last_node=node_list[-1]
    multi_level_dict[last_node]=val
def get_result_df(root_path,same_length=False,groupby="iterations",filter_list=[],only_mean=False,rerun=False,suffix="jjson"):
    results_path=os.path.join(root_path,"results.pickle")
    results_path_exist=glob.glob(results_path)
#     print(results_path)
    if results_path_exist and not rerun:
        results_path=results_path_exist[0]
        print("found saved results")
        with open(results_path, 'rb') as handle:
            result,result_mean = pickle.load(handle)
    else:
        path_sets=get_path_sets(root_path,same_length,suffix=suffix)
        # print(path_sets)
        result=AutoVivification()
        result_mean=AutoVivification()
        for path_set in path_sets:
            node=get_node(root_path,path_set)
            result_cur=merge_single_experiment_results(path_set,suffix=suffix)
            set_node_val(node,result,result_cur)
            result_merged_cur=result_cur.groupby(groupby).mean().reset_index()
            set_node_val(node,result_mean,result_merged_cur)
        with open(results_path, 'wb') as handle:
            pickle.dump([result,result_mean], handle, protocol=pickle.HIGHEST_PROTOCOL)
    if not only_mean:
        return result,result_mean
    else:
        return result
def extract_step_metric(result:dict,metric_name,step,data_name):
    result_return=[]
    for key, value in result.items():
        if  isinstance(value, pd.DataFrame) and metric_name in value:
            result_return.append([key,data_name,value[metric_name][step].tolist()])
        else:
            result_return.append([key,data_name,None])
    return result_return

def latex_two_f(x, y):                      # this is a demo function that takes in two ints and
    if x>1:
        x="{:#.4G}".format(x)
    else:
        x="{:.3f}".format(x)
    if y>1:
        y="{:#.4G}".format(y)
    else:
        y="{:.3f}".format(y)
    return "&"+str(x) + "$_{("+str(y)+")}$"        # concatenate them as str
vec_latex_two_f = np.vectorize(latex_two_f)
def latex_single_f_latex(x):                      # this is a demo function that takes in two ints and
    if x>1:
        x="{:#.4G}".format(x)
    else:
        x="{:.3f}".format(x)
    return "&"+str(x)        # concatenate them as str
vec_latex_single_f_latex = np.vectorize(latex_single_f_latex)
def latex_single_f(x):                      # this is a demo function that takes in two ints and
    if x>1:
        x="{:.3f}".format(x)
    else:
        x="{:.3f}".format(x)
    return str(x)        # concatenate them as str
vec_latex_single_f = np.vectorize(latex_single_f)

def to_mean(result_dataframe):
    mean=result_dataframe.applymap(func=np.mean)
    result=pd.DataFrame(vec_latex_single_f(mean),index=mean.index,columns=mean.columns)
    return result

def to_latex(result_dataframe):
    std=result_dataframe.applymap(func=np.std)
    mean=result_dataframe.applymap(func=np.mean)
    result=pd.DataFrame(vec_latex_single_f_latex(mean),index=mean.index,columns=mean.columns)
    result_std=pd.DataFrame(vec_latex_two_f(mean, std),index=mean.index,columns=mean.columns)
    return result,result_std
def get_freq_singe(x):
    if not isinstance(x, list):
        if pd.isnull(x):
            return 0
        else:
            return 1
    else:
        return len(x)
def get_round_value(x):
    if not isinstance(x, list):
        return 1
    else:
        return [round(i, 2) for i in x]
def to_round(result_dataframe):
    frequency=result_dataframe.applymap(func=get_round_value)
    return frequency
def to_freq(result_dataframe):
    frequency=result_dataframe.applymap(func=get_freq_singe)
    return frequency

import itertools

def plot_metrics(name_results_pair:dict,plots_y_partition:str="metrics_NDCG",errbar=True,
plots_x_partition:str="iterations",groupby="iterations",ax=None,graph_param={})->None:

    '''
        name_results_pair:{method_name:result_dataframe}
        plots_partition: key name in each result_dataframe which need to be plotted
    '''

    prop_cycle = plt.rcParams['axes.prop_cycle']
    colors_list = prop_cycle.by_key()['color']
    colors=itertools.cycle(colors_list)
    marker = itertools.cycle((',', '+', '.', 'o', '*'))
    if ax:
        plot=ax
    else:
        plot=plt
    for algo_name in name_results_pair:
            algo_result=name_results_pair[algo_name]
            print(type(algo_result))
            mean=algo_result.groupby(groupby).mean().reset_index()
            std=algo_result.groupby(groupby).std().reset_index()
            if plots_x_partition not in mean.keys() or plots_y_partition not in mean.keys() :
                continue
#             assert plots_y_partition in algo_result, algo_name+" doesn't contain the partition "+plots_y_partition
            ndata=len(mean[plots_x_partition])
            if not errbar:
                plot.plot(mean[plots_x_partition],mean[plots_y_partition], marker = next(marker),color=next(colors), label=algo_name, markevery=ndata//3)
            else:
                plot.errorbar(mean[plots_x_partition],mean[plots_y_partition], yerr=std[plots_y_partition], marker = next(marker),color=next(colors), label=algo_name, markevery=ndata//3)
    if ax is None:
        gca=plot.gca()
        gca.set(**graph_param)
        plot.legend()

def plot(name_results_pair:dict,errbar=True, ax=None,graph_param={},\
                    desiredGradFairColorDict={})->None:

    '''
        name_results_pair:{method_name:result_dataframe}
    '''

    prop_cycle = plt.rcParams['axes.prop_cycle']
    colors_list = prop_cycle.by_key()['color']
    colors=itertools.cycle(colors_list)
    marker = itertools.cycle((',', '+', '.', 'o', '*'))
    if ax:
        plot=ax
    else:
        plot=plt
    for algo_name in name_results_pair:
            x,y=name_results_pair[algo_name]
            y=np.array(y)
            if len(y.shape)==1:
                y=y[:,None]
            yMean=y.mean(axis=1)
            yStd=y.std(axis=1)
            ndata=len(x)
#             assert plots_y_partition in algo_result, algo_name+" doesn't contain the partition "+plots_y_partition

            if desiredGradFairColorDict is None:
                color=next(colors)
            else:
                color=desiredGradFairColorDict[algo_name]
            if not errbar:
                plot.plot(x,y, marker = next(marker),color=color, label=algo_name, markevery=ndata//3)
            else:
                plot.errorbar(x,yMean, yStd, marker = next(marker),color=color, label=algo_name, markevery=ndata//3)
    if ax is None:
        gca=plot.gca()
        gca.set(**graph_param)
        plot.legend()

def paramIterationPlot(result:dict,metrics_name,step,ax=None,xlim=None,savepath=None):
    """
    This function plot the performance along different parameters.
    """
    if ax:
        plot=ax
    else:
        plot=plt
    marker = itertools.cycle(('v', '^', "<",">","p",'x',"X","D", 'o', '*'))
    for key, value in result.items():
        result_list=extractResultWithParam(value,[metrics_name],step)
        print(key)
        param,result_metric=result_list[0],result_list[1]

        if len(result_metric)>1:
#                 print(result_metric[i])
            mean=np.array([np.mean(i)for i in result_metric])
            std=np.array([np.std(i)for i in result_metric])
            ind=np.arange(mean.shape[0])
            if xlim:
                ind=param<=xlim
#                     print(ind)
            ndata=len(ind)
            plot.errorbar(param[ind],mean[ind],yerr=std[ind],marker = next(marker),label=key, markevery=ndata//3)
        else:
            plot.plot(param,result_metric,marker = next(marker),label=key)
def TradeoffPlot(result:dict,metrics_pair,step,\
                    desiredGradFairColorDict={},ax=None,xlim=None,savepath=None):
    """
    This function plot the tradeoff performance.
    """
    if ax:
        plot=ax
    else:
        plot=plt
    marker = itertools.cycle(('v', '^', "<",">","p",'x',"X","D", 'o', '*'))
    for key, value in result.items():
        result_list=extractResultWithParam(value,metrics_pair,step)
        xResult=np.array([np.mean(i)for i in result_list[1]])
        yResult=np.array([np.mean(i)for i in result_list[2]])
        ndata=len(xResult)
        plot.plot(xResult,yResult,marker = next(marker),label=key, markevery=ndata//3)

def RequirementPlot(result:dict,metrics_pair,step,\
                    desiredGradFairColorDict={},ax=None,xlim=None,savepath=None):
    """
    This function plot the tradeoff performance.
    """
    if ax:
        plot=ax
    else:
        plot=plt
    prop_cycle = plt.rcParams['axes.prop_cycle']
    colors_list = prop_cycle.by_key()['color']
    colors=itertools.cycle(colors_list)
    marker = itertools.cycle(('v', '^', "<",">","p",'x',"X","D", 'o', '*'))
    for key, value in result.items():
        result_list=extractResultWithParam(value,metrics_pair,step)
        xResult=np.array([np.mean(i)for i in result_list[1]])
        yResult=np.array([np.mean(i)for i in result_list[2]])
        ndata=len(xResult)
        xSortedid=np.argsort(xResult)
        xResult=xResult[xSortedid]
        yResult=yResult[xSortedid]
        yResult=np.array([np.max(yResult[:i+1])for i in range(ndata)])
        color=next(colors)
        if key in desiredGradFairColorDict:
            color=desiredGradFairColorDict[key]
        plot.plot(xResult,yResult,color=color,marker = next(marker),label=key, markevery=ndata//3)

def TradeoffScatter(result:dict,metrics_pair,step,\
                    desiredGradFairColorDict={},ax=None,xlim=None,savepath=None):
    """
    This function plot the tradeoff performance.
    """
    if ax:
        plot=ax
    else:
        plot=plt
    marker = itertools.cycle(("o"))
    for key, value in result.items():
        result_list=extractResult(value,metrics_pair,step)
        xResult=np.mean(result_list[0])
        yResult=np.mean(result_list[1])
        # ndata=len(xResult)
        if key in desiredGradFairColorDict:
            color=desiredGradFairColorDict[key]
            plot.scatter(xResult,yResult,color=color,marker = next(marker),label=key,s=200)
        else:
            plot.scatter(xResult,yResult,marker = next(marker),label=key,s=200)
def extractResultWithParam(cur_res_split,res_names=[],step=0):
    tradeoff_params=cur_res_split.keys()
    tradeoff_nums=[]
    tradeoff_params_filtered=[]
    for i in tradeoff_params:
        cur_num=i.split("_")[-1]
        if cur_num.lstrip('-').replace('.','',1).isdigit():
             tradeoff_nums.append(float(cur_num))
             tradeoff_params_filtered.append(i)
    tradeoff_params=tradeoff_params_filtered
    tradeoff_nums=np.array(tradeoff_nums)
    argsort=np.argsort(tradeoff_nums)
    tradeoff_nums=tradeoff_nums[argsort]
    result=[tradeoff_nums]
    for res_name in res_names:
        res_cur_name=[cur_res_split[tradeoff_param][res_name][step].tolist() for tradeoff_param in tradeoff_params ]
        res_cur_name=np.array(res_cur_name)[argsort]
        result.append(res_cur_name)
    return result
def extractResult(cur_res_split,res_names=[],step=0):
    result=[]
    for res_name in res_names:
        res_cur_name=[cur_res_split[res_name][step].tolist()]
        res_cur_name=np.array(res_cur_name)
        result.append(res_cur_name)
    return result

def getGrandchildNode(InputDict,GrandchildNode):
    """
    This function get the grandchild node.
    """
    resDict={}
    for keyL1, valueL1 in InputDict.items():
            resDict[keyL1]=valueL1[GrandchildNode]
    return resDict

def filteroutNone(result_list):
    """
    This function filter out results list after we use function extract_step_metric
    """
    result_list_new=[]
    for result_list_cur in result_list:
        if result_list_cur[2] is not None:
            result_list_new.append(result_list_cur)
    return result_list_new
def iterate_applyfunction(result,fcn):
    """
    This function return iteratively apply function fcn to result if result is a dataframe.
    """
    for key in result:
        if isinstance(result[key], pd.DataFrame):
            fcn(result[key])
        else:
            iterate_applyfunction(result[key],fcn)

def  reorderDict(dictInput,keys):
    """
    This function reorder the key of an dict input.
    """
    dictOutput={}
    for key in keys:
        if key in dictInput:
            dictOutput[key]=dictInput[key]
    for key in dictInput:
        if key not in dictOutput:
            dictOutput[key]=dictInput[key]
    return  dictOutput

def reorder(desiredList,curList):
    """
    This function index mapping from the curList according to desiredList.
    """
    reoderIndex=[]
    desiredListCur=[]
    for legend in desiredList:
        if legend in curList:
            desiredListCur.append(legend)
    desiredList=desiredListCur
    for legend in curList:
        if legend not in desiredList:
            desiredList.append(legend)
    for legend in desiredList:
        if legend in curList:
            reoderIndex.append(curList.index(legend))
    return reoderIndex
def reorderLegend(desiredLegend,ax,returnHandles=False):
    """
    This function index reorder the legend to desiredLegend.
    """
    handles, labels = plt.gca().get_legend_handles_labels()
    order=reorder(desiredLegend,labels)
    handles=[handles[idx] for idx in order]
    labels=[labels[idx] for idx in order]
    legend=ax.legend(handles,labels)
    if returnHandles:
        return legend,handles,labels
    return legend
import matplotlib
def figureConfig():
    font = {'family' : 'normal',
            'size'   : 12}
    matplotlib.rc('font', **font)
def setScaleFunction(a,b=1,low=True):
    """
    This function returns the scale conversion function.
    """
    if low:
        forwad=lambda x: np.log(((x-a)*b))
        inverse=lambda x: a+(np.exp(x))/b
    else:
        forwad=lambda x: -np.log(((a-x)*b))
        inverse=lambda x: a-(np.exp(-x))/b
    return [forwad,inverse]
def export_legend(legend, filename="legend.png", expand=[-5,-5,5,5]):
    """
    This function plot and save the legend.
    """
    fig  = legend.figure
    fig.canvas.draw()
    bbox  = legend.get_window_extent()
    bbox = bbox.from_extents(*(bbox.extents + np.array(expand)))
    bbox = bbox.transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(filename, dpi=600,pad_inches=0.05, bbox_inches=bbox)

max_stars_repo_name: Taosheng-ty/BatchExpLaunch

max_issues_repo_name: Taosheng-ty/BatchExpLaunch

max_forks_repo_name: Taosheng-ty/BatchExpLaunch


text:
import sys
import os
import time

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm

from attacut import Tokenizer, __version__, preprocessing, utils, models

# from https://github.com/pytorch/pytorch/issues/1494#issuecomment-305993854
from multiprocessing import set_start_method
try:
    set_start_method('spawn')
except RuntimeError:
    pass

SEP = "|"

# For models tranied with weight-decay, some weights are close to zero causing denormal
# ref: https://en.wikipedia.org/wiki/Denormal_number
# This block is aken from https://github.com/pytorch/pytorch/issues/19651#issuecomment-486170718
if not torch.set_flush_denormal(True):
    print("Unable to set flush denormal")
    print("Pytorch compiled without advanced CPU")
    print("at: https://github.com/pytorch/pytorch/blob/84b275b70f73d5fd311f62614bccc405f3d5bfa3/aten/src/ATen/cpu/FlushDenormal.cpp#L13")

def get_argument(dict, name, default):
    v = dict.get(name)
    return v if v is not None else default

class AttaCutCLIDataset(Dataset):
      def __init__(self, src, tokenizer, device):
          self.src = src

          self.tokenizer = tokenizer

          self.device = device

          self.data = []

          self.total_lines = 0
          with open(self.src, "r") as fin:
              for i, txt in enumerate(fin):
                txt = preprocessing.TRAILING_SPACE_RX.sub("", txt)

                tokens, features = self.tokenizer.dataset.make_feature(txt)

                inputs = (
                    features,
                    torch.zeros(features[1])  # dummy label when won't need it here
                )

                (x, seq), labels, _= self.tokenizer.dataset.prepare_model_inputs(
                  inputs,
                )

                x = torch.squeeze(x)

                self.data.append((((x, seq), labels), tokens))
                self.total_lines += 1

      def __len__(self):
          return self.total_lines

      def __getitem__(self, idx):
          return self.data[idx]

def collate_fn(tokenizer, batch, device):
    inputs, tokens = [], []

    for x, t in batch:
      inputs.append(x)
      tokens.append(t)

    (x, seq), labels, perm_idx = tokenizer.dataset.collate_fn(inputs)

    return ((x.to(device), seq.to(device)), labels, perm_idx), tokens

def main(src, model, num_cores=4, batch_size=32, dest=None, device="cpu"):

    assert num_cores >= 0, "Input given to <num-thread> should greather than or equal one"

    if not src:
      print(__doc__)
      sys.exit(0)

    # for a custom model, use the last dir's name.
    model_name = model.split("/")[-1]

    dest = dest if dest is not None else utils.add_suffix_to_file_path(src, f"tokenized-by-{model_name}")

    print(f"Tokenizing {src}")
    print(f"Using {src}")
    print(f"Output: {dest}")

    tokenizer = Tokenizer(model)

    total_lines = utils.wc_l(src)

    if num_cores == 0:
        print(f"Use main process processing for {total_lines} lines")
    else:
        print(f"Use {num_cores} cores for processing for {total_lines} lines")

    print(f"device={device}")

    start_time = time.time()
    ds = AttaCutCLIDataset(src, tokenizer, device)
    dataloader = DataLoader(
      ds,
      batch_size=batch_size,
      shuffle=False,
      num_workers=0, # only use main process
      collate_fn=lambda batch: collate_fn(tokenizer, batch, device)
    )

    tokenizer.model.to(device)

    tokenizer.model.eval()
    results = []

    with torch.no_grad(), \
        tqdm(total=total_lines) as tq, \
        open(dest, "w") as fout:

            for batch in dataloader:
                (x, labels, perm_idx), tokens = batch
                logits = tokenizer.model(x)

                seq_lengths = x[1]

                preds = tokenizer.model.decode(logits, seq_lengths)

                perm_idx = perm_idx.cpu().detach()

                for ori_ix, after_sorting_ix in enumerate(np.argsort(perm_idx)):
                    pred = preds[after_sorting_ix]
                    token = tokens[ori_ix]

                    words = preprocessing.find_words_from_preds(token, pred)
                    fout.write("%s\n" % SEP.join(words))

                tq.update(n=seq_lengths.shape[0])

    time_took = time.time() - start_time

    return time_took