text:
import warnings

import numpy as np
import pandas as pd

from powersimdata.design.compare.generation import calculate_plant_difference
from powersimdata.design.compare.transmission import (
    calculate_branch_difference,
    calculate_dcline_difference,
)
from powersimdata.design.investment import const
from powersimdata.design.investment.create_mapping_files import (
    bus_to_neem_reg,
    bus_to_reeds_reg,
)
from powersimdata.design.investment.inflation import calculate_inflation
from powersimdata.input.check import _check_grid_models_match
from powersimdata.utility.distance import haversine

def merge_keep_index(df1, df2, **kwargs):
    """Execute a pandas DataFrame merge, preserving the index of the first dataframe.

    :param pandas.DataFrame df1: first data frame, to call pandas merge from.
    :param pandas.DataFrame df2: second data frame, argument to pandas merge.
    :param \\*\\*kwargs: arbitrary keyword arguments passed to pandas merge call.
    :return: (*pandas.DataFrame*) -- df1 merged with df2 with indices preserved.
    """
    return df1.reset_index().merge(df2, **kwargs).set_index(df1.index.names)

def append_keep_index_name(df1, other, *args, **kwargs):
    """Execute a pandas DataFrame append, preserving the index name of the dataframe.

    :param pandas.DataFrame df1: first data frame, to call pandas append from.
    :param pandas.DataFrame/pandas.Series/list: first argument to pandas append method.
    :param \\*args: arbitrary positional arguments passed to pandas append call.
    :param \\*\\*kwargs: arbitrary keyword arguments passed to pandas append call.
    :return: (*pandas.DataFrame*) -- df1 appended with other with index name preserved.
    """
    original_index_name = df1.index.name
    new_df = df1.append(other, *args, **kwargs)
    new_df.index.name = original_index_name
    return new_df

def calculate_ac_inv_costs(
    scenario,
    sum_results=True,
    exclude_branches=None,
    base_grid=None,
):
    """Calculate cost of upgrading AC lines and/or transformers in a scenario.
    NEEM regions are used to find regional multipliers.

    :param powersimdata.scenario.scenario.Scenario scenario: scenario instance.
    :param bool sum_results: whether to sum data frame for each branch type. Defaults to
        True.
    :param powersimdata.input.grid.Grid base_grid: a Grid to compare against. If None,
        the grid model and interconnect from the ``scenario`` are used to instantiate a
        corresponding unmodified Grid.
    :return: (*dict*) -- keys are {'line_cost', 'transformer_cost'}, values are either
        float if ``sum_results``, or pandas Series indexed by branch ID.
        Whether summed or not, values are $USD, inflation-adjusted to today.
    """
    grid_differences = scenario.get_grid()
    if base_grid is None:
        base_grid = scenario.get_base_grid()
    else:
        _check_grid_models_match(base_grid, grid_differences)

    # find upgraded AC lines
    capacity_difference = calculate_branch_difference(
        base_grid.branch, grid_differences.branch
    )
    grid_differences.branch = grid_differences.branch.assign(
        rateA=capacity_difference["diff"].to_numpy()
    )
    grid_differences.branch = grid_differences.branch.query("rateA != 0.0")
    if exclude_branches is not None:
        present_exclude_branches = set(exclude_branches) & set(
            grid_differences.branch.index
        )
        grid_differences.branch.drop(index=present_exclude_branches, inplace=True)

    costs = _calculate_ac_inv_costs(grid_differences, sum_results)
    return costs

def _calculate_ac_inv_costs(grid_new, sum_results=True):
    """Calculate cost of upgrading AC lines and/or transformers. NEEM regions are
    used to find regional multipliers. Note that a transformer winding is considered
    as a transformer.

    :param powersimdata.input.grid.Grid grid_new: grid instance.
    :param bool sum_results: whether to sum data frame for each branch type. Defaults to
        True.
    :return: (*dict*) -- keys are {'line_cost', 'transformer_cost'}, values are either
        float if ``sum_results``, or pandas Series indexed by branch ID.
        Whether summed or not, values are $USD, inflation-adjusted to today.
    """

    def select_mw(x, cost_df):
        """Determine the closest kV/MW combination for a single branch and return
        the corresponding cost (in $/MW-mi).

        :param pandas.Series x: data for a single branch
        :param pandas.DataFrame cost_df: data frame with *'kV'*, *'MW'*, *'costMWmi'*
            as columns
        :return: (*pandas.Series*) -- series of [*'MW'*, *'costMWmi'*] to be assigned
            to branch.
        """
        underground_regions = ("NEISO", "NYISO J-K")
        filtered_cost_df = cost_df.copy()
        # Unless we are entirely within an underground region, drop this cost class
        if not (x.from_region == x.to_region and x.from_region in underground_regions):
            filtered_cost_df = filtered_cost_df.query("kV != 345 or MW != 500")
        # select corresponding cost table of selected kV
        filtered_cost_df = filtered_cost_df[filtered_cost_df["kV"] == x.kV]
        # get rid of NaN values in this kV table
        filtered_cost_df = filtered_cost_df[~filtered_cost_df["MW"].isna()]
        # find closest MW & corresponding cost
        filtered_cost_df = filtered_cost_df.iloc[
            np.argmin(np.abs(filtered_cost_df["MW"] - x.rateA))
        ]
        return filtered_cost_df.loc[["MW", "costMWmi"]]

    def get_branch_mult(x, bus_reg, ac_reg_mult, branch_lookup_alerted=set()):
        """Determine the regional multiplier based on kV and power (closest).

        :param pandas.Series x: data for a single branch.
        :param pandas.DataFrame bus_reg: data frame with bus regions.
        :param pandas.DataFrame ac_reg_mult: data frame with regional multipliers.
        :param set branch_lookup_alerted: set of (voltage, region) tuples for which
            a message has already been printed that this lookup was not found.
        :return: (*float*) -- regional multiplier.
        """
        # Select the highest voltage for transformers (branch end voltages should match)
        max_kV = bus.loc[[x.from_bus_id, x.to_bus_id], "baseKV"].max()  # noqa: N806
        # Average the multipliers for branches (transformer regions should match)
        regions = (x.from_region, x.to_region)
        region_mults = ac_reg_mult.loc[ac_reg_mult.name_abbr.isin(regions)]
        region_mults = region_mults.groupby(["kV", "MW"]).mean().reset_index()

        mult_lookup_kV = region_mults.loc[  # noqa: N806
            (region_mults.kV - max_kV).abs().idxmin()
        ].kV
        region_kV_mults = region_mults[region_mults.kV == mult_lookup_kV]  # noqa: N806
        region_kV_mults = region_kV_mults.loc[  # noqa: N806
            ~region_kV_mults.mult.isnull()
        ]
        if len(region_kV_mults) == 0:
            mult = 1
            if (mult_lookup_kV, regions) not in branch_lookup_alerted:
                print(f"No multiplier for voltage {mult_lookup_kV} in {regions}")
                branch_lookup_alerted.add((mult_lookup_kV, regions))
        else:
            mult_lookup_MW = region_kV_mults.loc[  # noqa: N806
                (region_kV_mults.MW - x.rateA).abs().idxmin(), "MW"
            ]
            mult = (
                region_kV_mults.loc[region_kV_mults.MW == mult_lookup_MW].squeeze().mult
            )
        return mult

    # import data
    ac_cost = pd.DataFrame(const.ac_line_cost)
    ac_reg_mult = pd.read_csv(const.ac_reg_mult_path)
    ac_reg_mult = ac_reg_mult.melt(
        id_vars=["kV", "MW"], var_name="name_abbr", value_name="mult"
    )
    try:
        bus_reg = pd.read_csv(const.bus_neem_regions_path, index_col="bus_id")
    except FileNotFoundError:
        bus_reg = bus_to_neem_reg(grid_new.bus)
        bus_reg.sort_index().to_csv(const.bus_neem_regions_path)
    xfmr_cost = pd.read_csv(const.transformer_cost_path, index_col=0).fillna(0)
    xfmr_cost.columns = [int(c) for c in xfmr_cost.columns]
    # Mirror across diagonal
    xfmr_cost += xfmr_cost.to_numpy().T - np.diag(np.diag(xfmr_cost.to_numpy()))

    # check that all buses included in this file and lat/long values match,
    # otherwise re-run mapping script on mis-matching buses. These buses are missing
    # in region file
    bus = grid_new.bus
    mapped_buses = bus.query("index in @bus_reg.index")
    missing_bus_indices = set(bus.index) - set(bus_reg.index)
    mapped_buses = merge_keep_index(mapped_buses, bus_reg, how="left", on="bus_id")
    # these buses have incorrect lat/lon values in the region mapping file.
    #   re-running the region mapping script on those buses only.
    misaligned_bus_indices = mapped_buses[
        ~np.isclose(mapped_buses.lat_x, mapped_buses.lat_y)
        | ~np.isclose(mapped_buses.lon_x, mapped_buses.lon_y)
    ].index
    all_buses_to_fix = set(missing_bus_indices) | set(misaligned_bus_indices)
    # fix the identified buses, if necessary
    if len(all_buses_to_fix) > 0:
        bus_fix = bus_to_neem_reg(bus.query("index in @all_buses_to_fix"))
        fix_cols = ["name_abbr", "lat", "lon"]
        corrected_bus_mappings = bus_fix.loc[misaligned_bus_indices, fix_cols]
        new_bus_mappings = bus_fix.loc[missing_bus_indices, fix_cols]
        bus_reg.loc[misaligned_bus_indices, fix_cols] = corrected_bus_mappings
        bus_reg = append_keep_index_name(bus_reg, new_bus_mappings)

    bus_reg.drop(["lat", "lon"], axis=1, inplace=True)

    # Add extra information to branch data frame
    branch = grid_new.branch
    branch.loc[:, "kV"] = bus.loc[branch.from_bus_id, "baseKV"].tolist()
    branch.loc[:, "from_region"] = bus_reg.loc[branch.from_bus_id, "name_abbr"].tolist()
    branch.loc[:, "to_region"] = bus_reg.loc[branch.to_bus_id, "name_abbr"].tolist()
    # separate transformers and lines
    t_mask = branch["branch_device_type"].isin(["Transformer", "TransformerWinding"])
    transformers = branch[t_mask].copy()
    lines = branch[~t_mask].copy()
    if len(lines) > 0:
        # Find closest kV rating
        lines.loc[:, "kV"] = lines.apply(
            lambda x: ac_cost.loc[(ac_cost["kV"] - x.kV).abs().idxmin(), "kV"],
            axis=1,
        )
        lines[["MW", "costMWmi"]] = lines.apply(lambda x: select_mw(x, ac_cost), axis=1)

        lines["mult"] = lines.apply(
            lambda x: get_branch_mult(x, bus_reg, ac_reg_mult), axis=1
        )

        # calculate MWmi
        lines.loc[:, "lengthMi"] = lines.apply(
            lambda x: haversine((x.from_lat, x.from_lon), (x.to_lat, x.to_lon)), axis=1
        )
    else:
        new_columns = ["kV", "MW", "costMWmi", "mult", "lengthMi"]
        lines = lines.reindex(columns=[*lines.columns.tolist(), *new_columns])
    lines.loc[:, "MWmi"] = lines["lengthMi"] * lines["rateA"]

    # calculate cost of each line
    lines.loc[:, "cost"] = lines["MWmi"] * lines["costMWmi"] * lines["mult"]

    # calculate transformer costs
    if len(transformers) > 0:
        transformers["per_MW_cost"] = transformers.apply(
            lambda x: xfmr_cost.iloc[
                xfmr_cost.index.get_loc(
                    bus.loc[x.from_bus_id, "baseKV"], method="nearest"
                ),
                xfmr_cost.columns.get_loc(
                    bus.loc[x.to_bus_id, "baseKV"], method="nearest"
                ),
            ],
            axis=1,
        )
        transformers["mult"] = transformers.apply(
            lambda x: get_branch_mult(x, bus_reg, ac_reg_mult), axis=1
        )
    else:
        # Properly handle case with no transformers, where apply returns wrong dims
        transformers["per_MW_cost"] = []
        transformers["mult"] = []

    transformers["cost"] = (
        transformers["rateA"] * transformers["per_MW_cost"] * transformers["mult"]
    )

    lines.cost *= calculate_inflation(2010)
    transformers.cost *= calculate_inflation(2020)
    if sum_results:
        return {
            "line_cost": lines.cost.sum(),
            "transformer_cost": transformers.cost.sum(),
        }
    else:
        return {"line_cost": lines.cost, "transformer_cost": transformers.cost}

def calculate_dc_inv_costs(scenario, sum_results=True, base_grid=None):
    """Calculate cost of upgrading HVDC lines in a scenario.

    :param powersimdata.scenario.scenario.Scenario scenario: scenario instance.
    :param bool sum_results: whether to sum series to return total cost. Defaults to
        True.
    :param powersimdata.input.grid.Grid base_grid: a Grid to compare against. If None,
        the grid model and interconnect from the ``scenario`` are used to instantiate a
        corresponding unmodified Grid.
    :return: (*pandas.Series/float*) -- cost of upgrading HVDC lines, in $USD,
        inflation-adjusted to today. If ``sum_results``, a float is returned, otherwise
        a Series.
    """
    grid_differences = scenario.get_grid()
    if base_grid is None:
        base_grid = scenario.get_base_grid()
    else:
        _check_grid_models_match(base_grid, grid_differences)

    # find upgraded DC lines
    capacity_difference = calculate_dcline_difference(
        base_grid.dcline, grid_differences.dcline
    )
    grid_differences.dcline = grid_differences.dcline.assign(
        Pmax=capacity_difference["diff"].to_numpy()
    )
    grid_differences.dcline = grid_differences.dcline.query("Pmax != 0.0")

    costs = _calculate_dc_inv_costs(grid_differences, sum_results)
    return costs

def _calculate_dc_inv_costs(grid_new, sum_results=True):
    """Calculate cost of upgrading HVDC lines.

    :param powersimdata.input.grid.Grid grid_new: grid instance.
    :param bool sum_results: whether to sum series to return total cost. Defaults to
        True.
    :return: (*pandas.Series/float*) -- cost of upgrading HVDC lines, in $USD,
        inflation-adjusted to today. If ``sum_results``, a float is returned, otherwise
        a Series.
    """

    def _calculate_single_line_cost(line, bus):
        """Calculate cost of upgrading a single HVDC line.

        :param pandas.Series line: HVDC line series featuring *'from_bus_id'*',
            *'to_bus_id'* and *'Pmax'*.
        :param pandas.Dataframe bus: bus data frame featuring *'lat'*, *'lon'*.
        :return: (*float*) -- HVDC line upgrade cost in $2015.
        """
        # Calculate distance
        from_lat = bus.loc[line.from_bus_id, "lat"]
        from_lon = bus.loc[line.from_bus_id, "lon"]
        to_lat = bus.loc[line.to_bus_id, "lat"]
        to_lon = bus.loc[line.to_bus_id, "lon"]
        miles = haversine((from_lat, from_lon), (to_lat, to_lon))
        # Calculate cost
        total_cost = line.Pmax * (
            miles * const.hvdc_line_cost["costMWmi"] * calculate_inflation(2015)
            + 2 * const.hvdc_terminal_cost_per_MW * calculate_inflation(2020)
        )
        return total_cost

    bus = grid_new.bus
    dcline = grid_new.dcline

    # if any dclines, do calculations, otherwise, return 0 costs.
    if len(dcline != 0):
        dcline_costs = dcline.apply(_calculate_single_line_cost, args=(bus,), axis=1)
        if sum_results:
            return dcline_costs.sum()
        else:
            return dcline_costs
    else:
        return 0.0

def calculate_gen_inv_costs(
    scenario,
    year,
    cost_case,
    sum_results=True,
    base_grid=None,
):
    """Calculate cost of upgrading generators in a scenario. ReEDS regions are used to
    find regional multipliers.

    :param powersimdata.scenario.scenario.Scenario scenario: scenario instance.
    :param int/str year: building year.
    :param str cost_case: ATB cost case of data. *'Moderate'*: mid cost case,
        *'Conservative'*: generally higher costs, *'Advanced'*: generally lower costs
    :param bool sum_results: whether to sum data frame for plant costs. Defaults to
        True.
    :param powersimdata.input.grid.Grid base_grid: a Grid to compare against. If None,
        the grid model and interconnect from the ``scenario`` are used to instantiate a
        corresponding unmodified Grid.
    :return: (*pandas.Series*) -- Overnight generation investment cost.
        If ``sum_results``, indices are technologies and values are total cost.
        Otherwise, indices are IDs of plants (including storage, which is given
        pseudo-plant-IDs), and values are individual generator costs.
        Whether summed or not, values are $USD, inflation-adjusted to today.

    .. todo:: it currently uses one (arbitrary) sub-technology. The rest of the costs
        are dropped. Wind and solar will need to be fixed based on the resource supply
        curves.
    """
    grid_differences = scenario.get_grid()
    if base_grid is None:
        base_grid = scenario.get_base_grid()
    else:
        _check_grid_models_match(base_grid, grid_differences)

    # Find change in generation capacity
    capacity_difference = calculate_plant_difference(
        base_grid.plant, grid_differences.plant
    )
    grid_differences.plant = grid_differences.plant.assign(
        Pmax=capacity_difference["diff"].to_numpy()
    )
    grid_differences.plant = grid_differences.plant.query("Pmax >= 0.01")
    # Find change in storage capacity
    # Reindex so that we don't get NaN when calculating upgrades for new storage
    base_grid.storage["gen"] = base_grid.storage["gen"].reindex(
        grid_differences.storage["gen"].index, fill_value=0
    )
    grid_differences.storage["gen"].Pmax = (
        grid_differences.storage["gen"].Pmax - base_grid.storage["gen"].Pmax
    )
    grid_differences.storage["gen"]["type"] = "storage"

    costs = _calculate_gen_inv_costs(grid_differences, year, cost_case, sum_results)
    return costs

def _calculate_gen_inv_costs(grid_new, year, cost_case, sum_results=True):
    """Calculate cost of upgrading generators. ReEDS regions are used to find
    regional multipliers.

    :param powersimdata.input.grid.Grid grid_new: grid instance.
    :param int/str year: year of builds.
    :param str cost_case: ATB cost case of data. *'Moderate'*: mid cost case
        *'Conservative'*: generally higher costs, *'Advanced'*: generally lower costs.
    :raises ValueError: if year not 2020 - 2050, or cost case not an allowed option.
    :raises TypeError: if year not int/str or cost_case not str.
    :param bool sum_results: whether to sum data frame for plant costs. Defaults to
        True.
    :return: (*pandas.Series*) -- Overnight generation investment cost.
        If ``sum_results``, indices are technologies and values are total cost.
        Otherwise, indices are IDs of plants (including storage, which is given
        pseudo-plant-IDs), and values are individual generator costs.
        Whether summed or not, values are $USD, inflation-adjusted to today.

    .. note:: the function computes the total capital cost as:
        CAPEX_total = overnight CAPEX ($/MW) * Power capacity (MW) * regional multiplier
    """

    def load_cost(year, cost_case):
        """Load in base costs from NREL's 2020 ATB for generation technologies (CAPEX).

        :param int/str year: year of cost projections.
        :param str cost_case: ATB cost case of data (see
        :return: (*pandas.DataFrame*) -- cost by technology/subtype in $2018.

        .. todo:: it can be adapted in the future for FOM, VOM, & CAPEX. This data is
            pulled from the ATB xlsx file summary pages. Therefore, it currently uses
            default financials, but will want to create custom financial functions in
            the future.
        """
        cost = pd.read_csv(const.gen_inv_cost_path)
        cost = cost.dropna(axis=0, how="all")

        # drop non-useful columns
        cols_drop = cost.columns[
            ~cost.columns.isin(
                [str(x) for x in cost.columns[0:6]] + ["Metric", str(year)]
            )
        ]
        cost.drop(cols_drop, axis=1, inplace=True)

        # rename year of interest column
        cost.rename(columns={str(year): "value"}, inplace=True)

        # get rid of #refs
        cost.drop(cost[cost["value"] == "#REF!"].index, inplace=True)

        # get rid of $s, commas
        cost["value"] = cost["value"].str.replace("$", "", regex=True)
        cost["value"] = cost["value"].str.replace(",", "", regex=True).astype("float64")
        # scale from $/kW to $/MW
        cost["value"] *= 1000

        cost.rename(columns={"value": "CAPEX"}, inplace=True)

        # select scenario of interest
        if cost_case != "Moderate":
            # The 2020 ATB only has "Moderate" for nuclear, so we need to make due.
            warnings.warn(
                f"No cost data available for Nuclear for {cost_case} cost case, "
                "using Moderate cost case data instead"
            )
            new_nuclear = cost.query(
                "Technology == 'Nuclear' and CostCase == 'Moderate'"
            ).copy()
            new_nuclear.CostCase = cost_case
            cost = cost.append(new_nuclear, ignore_index=True)
        cost = cost[cost["CostCase"] == cost_case]
        cost.drop(["CostCase"], axis=1, inplace=True)

        return cost

    if isinstance(year, (int, str)):
        year = int(year)
        if year not in range(2020, 2051):
            raise ValueError("year not in range.")
    else:
        raise TypeError("year must be int or str.")

    if isinstance(cost_case, str):
        if cost_case not in ["Moderate", "Conservative", "Advanced"]:
            raise ValueError("cost_case not Moderate, Conservative, or Advanced")
    else:
        raise TypeError("cost_case must be str.")

    storage_plants = grid_new.storage["gen"].set_index(
        grid_new.storage["StorageData"].UnitIdx.astype(int)
    )
    plants = append_keep_index_name(grid_new.plant, storage_plants)
    plants = plants[
        ~plants.type.isin(["dfo", "other"])
    ]  # drop these technologies, no cost data

    # BASE TECHNOLOGY COST

    # load in investment costs $/MW
    gen_costs = load_cost(year, cost_case)
    # keep only certain (arbitrary) subclasses for now
    gen_costs = gen_costs[
        gen_costs["TechDetail"].isin(const.gen_inv_cost_techdetails_to_keep)
    ]
    # rename techs to match grid object
    gen_costs.replace(const.gen_inv_cost_translation, inplace=True)
    gen_costs.drop(["Key", "FinancialCase", "CRPYears"], axis=1, inplace=True)
    # ATB technology costs merge
    plants = merge_keep_index(
        plants, gen_costs, right_on="Technology", left_on="type", how="left"
    )

    # REGIONAL COST MULTIPLIER

    # Find ReEDS regions of plants (for regional cost multipliers)
    plant_buses = plants.bus_id.unique()
    try:
        bus_reg = pd.read_csv(const.bus_reeds_regions_path, index_col="bus_id")
        if not set(plant_buses) <= set(bus_reg.index):
            missing_buses = set(plant_buses) - set(bus_reg.index)
            bus_reg = bus_reg.append(bus_to_reeds_reg(grid_new.bus.loc[missing_buses]))
            bus_reg.sort_index().to_csv(const.bus_reeds_regions_path)
    except FileNotFoundError:
        bus_reg = bus_to_reeds_reg(grid_new.bus.loc[plant_buses])
        bus_reg.sort_index().to_csv(const.bus_reeds_regions_path)
    plants = merge_keep_index(
        plants, bus_reg, left_on="bus_id", right_index=True, how="left"
    )

    # Determine one region 'r' for each plant, based on one of two mappings
    plants.loc[:, "r"] = ""
    # Some types get regional multipliers via 'wind regions' ('rs')
    wind_region_mask = plants["type"].isin(const.regional_multiplier_wind_region_types)
    plants.loc[wind_region_mask, "r"] = plants.loc[wind_region_mask, "rs"]
    # Other types get regional multipliers via 'BA regions' ('rb')
    ba_region_mask = plants["type"].isin(const.regional_multiplier_ba_region_types)
    plants.loc[ba_region_mask, "r"] = plants.loc[ba_region_mask, "rb"]
    plants.drop(["rs", "rb"], axis=1, inplace=True)

    # merge regional multipliers with plants
    region_multiplier = pd.read_csv(const.regional_multiplier_path)
    region_multiplier.replace(const.regional_multiplier_gen_translation, inplace=True)
    plants = merge_keep_index(
        plants,
        region_multiplier,
        left_on=["r", "Technology"],
        right_on=["r", "i"],
        how="left",
    )

    # multiply all together to get summed CAPEX ($)
    plants.loc[:, "cost"] = (
        plants["CAPEX"] * plants["Pmax"] * plants["reg_cap_cost_mult"]
    )

    # sum cost by technology
    plants.loc[:, "cost"] *= calculate_inflation(2018)
    if sum_results:
        return plants.groupby(["Technology"])["cost"].sum()
    else:
        return plants["cost"]

max_stars_repo_name: c-voegele/PowerSimData

max_issues_repo_name: c-voegele/PowerSimData

max_forks_repo_name: c-voegele/PowerSimData


text:
# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
#
#   See COPYING file distributed along with the NiBabel package for the
#   copyright and license terms.
#
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
""" Preliminary MINC2 support

Use with care; I haven't tested this against a wide range of MINC files.

If you have a file that isn't read correctly, please send an example.

Test reading with something like::

    import nibabel as nib
    img = nib.load('my_funny.mnc')
    data = img.get_fdata()
    print(data.mean())
    print(data.max())
    print(data.min())

and compare against command line output of::

    mincstats my_funny.mnc
"""
import numpy as np

from .minc1 import Minc1File, MincHeader, Minc1Image, MincError

class Hdf5Bunch(object):
    """ Make object for accessing attributes of variable
    """

    def __init__(self, var):
        for name, value in var.attrs.items():
            setattr(self, name, value)

class Minc2File(Minc1File):
    """ Class to wrap MINC2 format file

    Although it has some of the same methods as a ``Header``, we use
    this only when reading a MINC2 file, to pull out useful header
    information, and for the method of reading the data out
    """

    def __init__(self, mincfile):
        self._mincfile = mincfile
        minc_part = mincfile['minc-2.0']
        # The whole image is the first of the entries in 'image'
        image = minc_part['image']['0']
        self._image = image['image']
        self._dim_names = self._get_dimensions(self._image)
        dimensions = minc_part['dimensions']
        self._dims = [Hdf5Bunch(dimensions[s]) for s in self._dim_names]
        # We don't currently support irregular spacing
        # https://en.wikibooks.org/wiki/MINC/Reference/MINC2.0_File_Format_Reference#Dimension_variable_attributes
        for dim in self._dims:
            if dim.spacing != b'regular__':
                raise ValueError('Irregular spacing not supported')
        self._spatial_dims = [name for name in self._dim_names
                              if name.endswith('space')]
        self._image_max = image['image-max']
        self._image_min = image['image-min']

    def _get_dimensions(self, var):
        # Dimensions for a particular variable
        # Differs for MINC1 and MINC2 - see:
        # https://en.wikibooks.org/wiki/MINC/Reference/MINC2.0_File_Format_Reference#Associating_HDF5_dataspaces_with_MINC_dimensions
        try:
            dimorder = var.attrs['dimorder'].decode()
        except KeyError:  # No specified dimensions
            return []
        # The dimension name list must contain only as many entries
        # as the variable has dimensions. This reduces errors when an
        # unnecessary dimorder attribute is left behind.
        return dimorder.split(',')[:len(var.shape)]

    def get_data_dtype(self):
        return self._image.dtype

    def get_data_shape(self):
        return self._image.shape

    def _get_valid_range(self):
        """ Return valid range for image data

        The valid range can come from the image 'valid_range' or
        failing that, from the data type range
        """
        ddt = self.get_data_dtype()
        info = np.iinfo(ddt.type)
        try:
            valid_range = self._image.attrs['valid_range']
        except (AttributeError, KeyError):
            valid_range = [info.min, info.max]
        else:
            if valid_range[0] < info.min or valid_range[1] > info.max:
                raise ValueError('Valid range outside input '
                                 'data type range')
        return np.asarray(valid_range, dtype=np.float64)

    def _get_scalar(self, var):
        """ Get scalar value from HDF5 scalar """
        return var[()]

    def _get_array(self, var):
        """ Get array from HDF5 array """
        return np.asanyarray(var)

    def get_scaled_data(self, sliceobj=()):
        """ Return scaled data for slice definition `sliceobj`

        Parameters
        ----------
        sliceobj : tuple, optional
            slice definition. If not specified, return whole array

        Returns
        -------
        scaled_arr : array
            array from minc file with scaling applied
        """
        if sliceobj == ():
            raw_data = np.asanyarray(self._image)
        else:  # Try slicing into the HDF array (maybe it's possible)
            try:
                raw_data = self._image[sliceobj]
            except (ValueError, TypeError):
                raw_data = np.asanyarray(self._image)[sliceobj]
            else:
                raw_data = np.asanyarray(raw_data)
        return self._normalize(raw_data, sliceobj)

class Minc2Header(MincHeader):

    @classmethod
    def may_contain_header(klass, binaryblock):
        return binaryblock[:4] == b'\211HDF'

class Minc2Image(Minc1Image):
    """ Class for MINC2 images

    The MINC2 image class uses the default header type, rather than a
    specific MINC header type - and reads the relevant information from
    the MINC file on load.
    """
    # MINC2 does not do compressed whole files
    _compressed_suffixes = ()
    header_class = Minc2Header

    @classmethod
    def from_file_map(klass, file_map, *, mmap=True, keep_file_open=None):
        # Import of h5py might take awhile for MPI-enabled builds
        # So we are importing it here "on demand"
        from ._h5py_compat import h5py
        holder = file_map['image']
        if holder.filename is None:
            raise MincError('MINC2 needs filename for load')
        minc_file = Minc2File(h5py.File(holder.filename, 'r'))
        affine = minc_file.get_affine()
        if affine.shape != (4, 4):
            raise MincError('Image does not have 3 spatial dimensions')
        data_dtype = minc_file.get_data_dtype()
        shape = minc_file.get_data_shape()
        zooms = minc_file.get_zooms()
        header = klass.header_class(data_dtype, shape, zooms)
        data = klass.ImageArrayProxy(minc_file)
        return klass(data, affine, header, extra=None, file_map=file_map)

load = Minc2Image.load

ext: py

reason: import numpy


text:
"""
mmirs_qrp
=========

Python script that quickly reduces longslit and MOS spectra from MMT/MMIRS
"""

import sys, os

#Py3 compatibility
py_vers = sys.version_info.major

from os.path import exists
from astropy.io import ascii as asc
from astropy.io import fits

from scipy.ndimage.interpolation import shift

import numpy as np

import matplotlib.pyplot as plt
import glob

from astropy.table import Table
from astropy import log

# Mod on 07/11/2017
from astropy.stats import sigma_clip
#from ccdproc import cosmicray_median

from scipy.optimize import curve_fit

# + on 20/11/2017
from matplotlib.backends.backend_pdf import PdfPages
from pylab import subplots_adjust

# + on 12/11/2017
import astropy.units as u

# + on 17/03/2018
import logging
formatter = logging.Formatter('%(asctime)s - %(module)12s.%(funcName)20s - %(levelname)s: %(message)s')
sh = logging.StreamHandler(sys.stdout)
sh.setLevel(logging.INFO)
sh.setFormatter(formatter)

pscale = 0.2012008872545049 # arcsec/pix

bbox_props = dict(boxstyle="square,pad=0.15", fc="w", alpha=0.75, ec="none")

class mlog:
    '''
    Main class to log information to stdout and ASCII file

    To execute:
    mylog = mlog(rawdir)._get_logger()

    Parameters
    ----------
    rawdir : str
      Full path for where raw files are

    Returns
    -------

    Notes
    -----
    Created by Chun Ly, 17 March 2018
     - Identical to logging function in MMTtools.mmirs_pipeline_taskfile
    '''

    def __init__(self,rawdir):
        self.LOG_FILENAME = rawdir + 'mmirs_qrp.log'
        self._log = self._get_logger()

    def _get_logger(self):
        loglevel = logging.INFO
        log = logging.getLogger(self.LOG_FILENAME) # + Mod on 14/12/2017
        if not getattr(log, 'handler_set', None):
            log.setLevel(logging.INFO)
            sh = logging.StreamHandler()
            sh.setFormatter(formatter)
            log.addHandler(sh)

            fh = logging.FileHandler(self.LOG_FILENAME)
            fh.setLevel(logging.INFO)
            fh.setFormatter(formatter)
            log.addHandler(fh)

            log.setLevel(loglevel)
            log.handler_set = True
        return log
#enddef

def gauss1d(x, a0, a, x0, sigma):
    return a0 + a * np.exp(-(x - x0)**2 / (2 * sigma**2))

def main(rawdir, prefix, bright=False, dither='ABApBp', flats=[],
         max_restrict=False):

    '''
    Main function of mmirs_qrp

    Parameters
    ----------
    rawdir : str
      Full path to MMIRS files

    prefix : str
      Prefix for files to process.  Code will search for dcorr files
      in rawdir + prefix + "*dcorr.fits*"
      For example, if your files are <target>_longslit.????.fits, provide 'target_longslit'

    bright : boolean
      Indicate whether a bright star/target is in slit.  If True,
      code will use bright star to determine offsets to shift spectra.
      If False, will use FITS header to determine dithered size
      (to be implemented). Default: False

    dither : str
      Dither sequence type.  Accepts 'ABApBp', 'ABAB' (to be implemented),
      and 'ABBA' (to be implemented). Default: 'ABApBp'

    flats: list
      List of files or seqno for flats.
      For example, flats=[1100,1101] or flats=['flat.1100','flat.1101']
      If not provided, flat fielding will not be performed

      NOTE: THERE ARE SOME BUGS WITH FLATFIELDING

    max_restrict : boolean
      Restrict search window for peak emission line when bright == True.
      Default will search the entire data.

    Returns
    -------

    Notes
    -----
    Created by Chun Ly, 13 October 2017
     - Handle gzip files as inputs
     - Change data cube dimensions
     - Handle FITS overwrite
     - Added dithering keyword
     - Handle ABA'B' dithering to get "sky" frame
     - AB subtraction for background removal
     - Add difference data cube for combine
     - Compute average from background-subtracted, image-shifted images
    Modified by Chun Ly, 14 October 2017
     - Bug in shift call. Need to specify as row,column shift values
    Modified by Chun Ly, 16 October 2017
     - Documentation added
    Modified by Chun Ly, 19 October 2017
     - Attempt CR rejection with ccdproc.cosmicray_median
    Modified by Chun Ly, 20 October 2017
     - Handle ABBA dithering to get "sky" frame
    Modified by Chun Ly,  7 November 2017
     - Use astropy.stats.sigma_clip to mask CRs
    Modified by Chun Ly, 12 November 2017
     - Get FITS header dither values, compute differences from first frame
     - Write ASCII table with dither offsets
     - Handle dithering for bright == False using FITS dither info
     - Handle output ASCII table for bright == False

     - Quality Assurance: Compute and plot FWHM
     - Plot seqno on x-axis for FWHM plot
     - Aesthetics for plots
     - Require bright=True for FWHM calculations
    Modified by Chun Ly, 17 November 2017
     - Write npz file of arrays to expedite analysis
     - Write npz file in compressed form
     - Handle masked arrays
     - Simplify npz to limit to just masked arrays
     - Read in npz file, Use npz file for sigma_clip mask
     - Bug fix: seqno handling of .gz files
    Modified by Chun Ly, 18 November 2017
     - Minor bug fix: NAXIS2 -> NAXIS1 and vice versa
    Modified by Chun Ly, 20 November 2017
     - Plot fit to PSF profiles
    Modified by Chun Ly, 25 November 2017
     - PSF profiles showed double peak with a lower peak that is 20% of max,
       but visual inspection does not indicate double-peaked PSF.
     - Unclear of the full cause, but ultimately computed FWHM in the central
       200 pix without masking (masking seems to reject pixels from bright objects)
     - Masking was part of the cause as it would produce a double peak
       distribution
    Modified by Chun Ly, 26 November 2017
     - Run curve_fit to get accurate line center for shifting
     - Set default shifting as integer pixels since CRs are present
       in stack (may need to grow mask)
    Modified by Chun Ly, 9 December 2017
     - Compute transparency: Integrate flux for bright source and plot
     - Normalize spectrum by exposure time for proper transparency computation
     - Plotting aesthetic improvements: legend, labels
     - Normalize transparency value to best, plotting aesthetics
     - Plotting aesthetic improvements: different linestyle and widths, smaller
       legend

     - Incorporate flat fielding (still some unexpected problems)

    Modified by Chun Ly, 11 December 2017
     - Expand dither_tab stdout,
     - Add annotation of avg/median/sigma for FWHM
     - Handle unused FWHM subplots
     - Minor stdout for avg/median/sigma for FWHM
     - Annotation for source info in FWHM and transparency plots

    Modified by Chun Ly, 18 March 2018
     - Implement stdout and ASCII logging with mlog()
     - Remove silent and verbose boolean keyword

    Modified by Chun Ly, 27 November 2018
     - Add max_restrict for restricting window to search for bright star

    Modified by Chun Ly, 18 December 2018
     - Do weighted stack with 1/FWHM^2
    '''

    mylog = mlog(rawdir)._get_logger() # + on 18/03/2018

    mylog.info('Begin main ! ') # Mod on 18/03/2018

    if rawdir[-1] != '/': rawdir = rawdir + '/'

    dcorr_files = glob.glob(rawdir+prefix+'*dcorr.fits*')
    n_files     = len(dcorr_files)

    if n_files == 0:
        mylog.warn("!!! d-corr files not found. Using raw data !!!")
        dcorr_yes = 0

        raw_files = glob.glob(rawdir+prefix+'.????.fits*')
        n_files   = len(raw_files)

        hdu0 = fits.getheader(raw_files[0], ext=1)
    else:
        dcorr_yes = 1

        hdu0 = fits.getheader(dcorr_files[0])

    naxis1, naxis2 = hdu0['NAXIS1'], hdu0['NAXIS2']

    npz_file = rawdir+prefix+'.npz'
    if exists(npz_file):
        mylog.info('Reading : '+npz_file) # Mod on 18/03/2018
        npz0 = np.load(npz_file)

    d_cube0     = np.zeros((n_files, naxis2, naxis1))
    peak_val    = np.zeros(n_files)
    peak_val0   = np.zeros(n_files) # + on 26/11/2017
    shift_cube0 = np.zeros((n_files, naxis2, naxis1))

    exptime0 = np.zeros(n_files) # + on 09/12/2017

    # Define and combine flat files | + on 09/12/2017
    do_flat = 0
    if len(flats) != 0:
        if 'flat' in str(flats[0]):
            flats = [flat+'_dcorr.fits' for flat in flats]
        else:
            flats = ['flat.%04i_dcorr.fits' % t_seq for t_seq in flats]

        flat_arr = np.zeros((len(flats), naxis2, naxis1))
        for ff in range(len(flats)):
            t_flat, flat_hdr = fits.getdata(rawdir+flats[ff], header=True)

            flat_arr[ff] = t_flat / np.max(t_flat)

        flat_arr_mask = sigma_clip(flat_arr, sigma=3., iters=3, axis=0)

        flat_avg = np.ma.average(flat_arr_mask, axis=0)
        flat0 = flat_avg.data

        bad = np.where((np.isfinite(flat0) == False) | (flat0 == 0) |
                       (flat0 == np.min(flat0)))
        if len(bad) > 0: flat0[bad] = 1.0

        out_fits_file = rawdir+prefix+'_flat.fits'
        mylog.info('Writing : '+out_fits_file) # Mod on 18/03/2018
        fits.writeto(out_fits_file, flat0, flat_hdr, overwrite=True)
        do_flat = 1

    # Set this to use curvefit to compute fractional offset
    do_curvefit_center = 0 # + on 26/11/2017

    # Mod on 12/11/2017
    dither_az   = np.zeros(n_files)
    dither_el   = np.zeros(n_files)
    diff_cube0  = np.zeros((n_files, naxis2, naxis1))

    if dither == 'ABApBp' or dither == 'ABBA':
        i_off = [1, -1] * np.int(n_files/2)
        if n_files % 2 == 1: i_off.append(-1) # Odd number correction
        i_sky = np.arange(n_files)+np.array(i_off)
    mylog.info("i_sky : " + ", ".join(["%i" % sky for sky in i_sky]))

    for ii in range(n_files):
        if dcorr_yes:
            d_data, t_hdr = fits.getdata(dcorr_files[ii], header=True)
        else:
            d_data0, t_hdr = fits.getdata(raw_files[ii], ext=1, header=True) #last
            d_data1 = fits.getdata(raw_files[ii], ext=2) #first
            d_data  = np.float_(d_data0) - d_data1

        if do_flat: d_data = d_data / flat0

        d_cube0[ii] = d_data

        # Mod on 12/11/2017
        dither_az[ii] = t_hdr['INSTAZ']
        dither_el[ii] = t_hdr['INSTEL']

        exptime0[ii] = t_hdr['EXPTIME'] # + on 09/12/2017

        ## + on 19/10/2017. Mod on 07/11/2017
        #data_crfree, crmask = cosmicray_median(d_cube0[ii], thresh=5, rbox=11)
        #d_cube0[ii] = data_crfree

    #Flag pixels that are outliers | Mod on 07/11/2017
    #arr_med = np.median(d_cube0, axis=0)
    #arr_std = np.std(d_cube0, axis=0)
    #
    #d_cube_mask = np.zeros_like(d_cube0)
    #
    #for ii in range(n_files):
    #    mask = np.where(np.absolute((d_cube0[ii] - arr_med)/arr_std) >= 4)
    #    d_cube_mask[ii][mask] = 1

    # Mod on 19/10/2017
    for ii in range(n_files):
        t_sky  = d_cube0[i_sky[ii]]
        t_diff = d_cube0[ii] - t_sky

        # Mod on 07/11/2017
        #t_sky  = np.ma.masked_array(d_cube0[i_sky[ii]],
        #                            mask=d_cube_mask[i_sky[ii]])
        #t_diff = np.ma.masked_array(d_cube0[ii], mask=d_cube_mask[ii]) - t_sky
        diff_cube0[ii] = t_diff

        # Compute offsets using bright source
        if bright == True:
            med0_row = np.median(t_diff, axis=1)

            resize   = np.repeat(med0_row, naxis1).reshape((naxis2,naxis1))
            im_test  = t_diff - resize

            med0_col     = np.median(im_test, axis=0)

            if not max_restrict:
                peak_val[ii] = np.argmax(med0_col)
            else:
                idx = range(450,1000)
                peak_val[ii] = idx[0]+np.argmax(med0_col[idx])

            # + on 26/11/2017
            if do_curvefit_center:
                p0 = [0.0, np.max(med0_col), peak_val[ii], 2.0]
                x0 = np.arange(len(med0_col))
                popt, pcov    = curve_fit(gauss1d, x0, med0_col, p0=p0)
                peak_val0[ii] = popt[2]
        else:
            if ii == 0: mylog.info('Using FITS dither information')

    dither_diff = (dither_el[0] - dither_el) / pscale * u.pix

    # Mod on 12/11/2017
    if bright == True:
        if do_curvefit_center: # Mod on 26/11/2017
            shift_val = peak_val0[0] - peak_val0 # Mod on 26/11/2017
        else:
            shift_val = peak_val[0] - peak_val
    else:
        shift_val = dither_diff.value

    mylog.info('Shift values for spectra : ') # Mod on 18/03/2018
    # Mod on 12/11/2017
    if dcorr_yes:
        files_short = [str0.replace(rawdir,'') for str0 in dcorr_files]
    else:
        files_short = [str0.replace(rawdir,'') for str0 in raw_files]

    if bright == True:
        diff0  = dither_diff - shift_val * u.pix
        arr0   = [files_short, dither_az * u.arcsec, dither_el * u.arcsec,
                  dither_diff, shift_val * u.pix, diff0]
        names0 = ('file','dither_az','dither_el','dither_diff','shift_val',
                  'difference')
    else:
        arr0   = [files_short, dither_az * u.arcsec, dither_el * u.arcsec,
                  dither_diff]
        names0 = ('file','dither_az','dither_el','dither_diff')

    dither_tab = Table(arr0, names=names0)
    dither_tab.pprint(max_lines=-1, max_width=-1)

    # + on 12/11/2017

    try:
        import yaml
        yaml_pass = 1
    except ModuleNotFoundError or ImportError:
        mylog.info("Failed yaml import")
        yaml_pass = 0

    if yaml_pass:
        out_dither_file1 = rawdir+prefix+'_dither_ecsv.cat'
        mylog.info('Writing : '+out_dither_file1) # Mod on 18/03/2018
        dither_tab.write(out_dither_file1, format='ascii.ecsv', overwrite=True)

    # + on 12/11/2017
    out_dither_file2 = rawdir+prefix+'_dither.cat'
    mylog.info('Writing : '+out_dither_file2) # Mod on 18/03/2018
    dither_tab.write(out_dither_file2, format='ascii.fixed_width_two_line', overwrite=True)
    #log.info('### '+" ".join([str(a) for a in shift_val]))

    for ii in range(n_files):
        # Bug fix - Mod on 14/10/2017
        shift_cube0[ii] = shift(diff_cube0[ii], [0,shift_val[ii]])

    # + on 07/11/2017
    if not exists(npz_file):
        shift_cube0_mask = sigma_clip(shift_cube0, sigma=3., iters=5, axis=0)
    else:
        shift_cube0_mask = np.ma.array(shift_cube0, mask=npz0['shift_cube0_mask'])

    stack0 = np.ma.average(shift_cube0_mask, axis=0)

    # + on 17/11/2017
    mylog.info('Writing : '+npz_file) # Mod on 18/03/2018
    np.savez_compressed(npz_file, dither_tab=dither_tab,
                        shift_cube0_mask=shift_cube0_mask.mask)
    #stack0d=stack0.data, stack0m=stack0.mask,

    # Mod on 07/11/2017
    fits.writeto(rawdir+prefix+'_stack.fits', stack0.data, overwrite=True)
    #fits.writeto(rawdir+prefix+'_stack.fits', stack0, overwrite=True)

    if bright == True:
        # Compute FWHM and plot | + on 12/11/2017
        out_fwhm_pdf = rawdir+prefix+'_stack_FWHM.pdf'

        # + on 20/11/2017
        pp = PdfPages(out_fwhm_pdf)
        ncols, nrows = 3, 2

        FWHM0 = np.zeros(n_files)
        if dcorr_yes:
            seqno = [str0.replace('.gz','').replace('_dcorr.fits','')[-4:] for \
                     str0 in dcorr_files] # Mod on 17/11/2017
        else:
            seqno = [str0.replace('.gz','').replace('.fits','')[-4:] for \
                     str0 in raw_files]

        for ii in range(n_files):
            if ii % (ncols*nrows) == 0: # + on 20/11/2017
                fig, ax = plt.subplots(nrows, ncols)

            row, col = ii / ncols % nrows, ii % ncols # + on 20/11/2017

            im0    = shift_cube0_mask[ii].data
            med0   = np.median(im0[1024-100:1024+100], axis=0) # Mod on 25/11/2017
            x0     = np.arange(len(med0))
            x0_max = np.argmax(med0)
            y0_max = np.max(med0)

            p0         = [0.0, y0_max, x0_max, 2.0]

            popt, pcov = curve_fit(gauss1d, x0, med0, p0=p0)
            FWHM0[ii]  = popt[3]*2*np.sqrt(2*np.log(2)) * pscale

            # + on 20/11/2017
            ax[row,col].plot((x0-popt[2])*pscale, med0/y0_max, color='black')
            ax[row,col].plot((x0-popt[2])*pscale, gauss1d(x0, *popt)/y0_max,
                             color='blue', alpha=0.5)

            # + on 20/11/2017
            if row == nrows-1:
                ax[row,col].set_xlabel('X [arcsec]')
            else:
                if ((n_files-1)-ii) > ncols-1:
                    ax[row,col].set_xticklabels([])

            # + on 11/12/2017
            if ii == n_files-1:
                for cc in range(ncols): ax[row,cc].set_xlabel('X [arcsec]')

            # + on 20/11/2017
            if col == 0:
                ax[row,col].set_ylabel('Normalized Flux')
            else: ax[row,col].set_yticklabels([])

            # + on 20/11/2017
            ax[row,col].annotate(seqno[ii], [0.025,0.975], ha='left',
                                 va='top', xycoords='axes fraction',
                                 weight='bold', fontsize=10)
            ax[row,col].annotate('FWHM = %.2f"' % FWHM0[ii], [0.975,0.975],
                                 ha='right', va='top', xycoords='axes fraction',
                                 weight='bold', fontsize=10)

            # + on 20/11/2017
            ax[row,col].set_xlim([-2.5,2.5])
            ax[row,col].set_ylim([-0.05,1.1])

            # + on 11/12/2017
            if ii == n_files-1:
                for cc in range(col+1,ncols): ax[row,cc].axis('off')
                for rr in range(row+1,nrows):
                    for cc in range(ncols): ax[rr,cc].axis('off')

            # + on 20/11/2017
            if ii % (ncols*nrows) == ncols*nrows-1 or ii == n_files-1:
                subplots_adjust(left=0.025, bottom=0.025, top=0.975,
                                right=0.975, wspace=0.02, hspace=0.02)
                fig.set_size_inches(8,6)
                fig.savefig(pp, format='pdf', bbox_inches='tight')
        #endfor

        fig, ax = plt.subplots() # + on 20/11/2017
        ax.plot(seqno, FWHM0, marker='o', color='b', alpha=0.5)
        ax.set_xlabel('Image Frame No.')
        ax.set_ylabel('FWHM [arcsec]')
        ax.minorticks_on()

        filt0, disp0, ap0 = t_hdr['FILTER'], t_hdr['DISPERSE'], \
                            t_hdr['APERTURE']

        label0 = '%s   %s   %s   %s' % (prefix, filt0, disp0, ap0)
        ax.annotate(label0, [0.025,0.975], ha='left', va='top', weight='bold',
                    xycoords='axes fraction', fontsize=11, bbox=bbox_props)

        # + on 11/12/2017
        avg_FWHM0 = np.average(FWHM0)
        med_FWHM0 = np.median(FWHM0)
        sig_FWHM0 = np.std(FWHM0)
        txt0  = 'Average : %.2f"\n Median : %.2f"\n' % (avg_FWHM0, med_FWHM0)
        txt0 += r'$\sigma$ : %.2f"' % sig_FWHM0
        # Mod on 18/03/2018
        mylog.info('Average : %.2f"' % avg_FWHM0)
        mylog.info('Median  : %.2f"' % med_FWHM0)
        mylog.info('Sigma   : %.2f"' % sig_FWHM0)

        ax.annotate(txt0, [0.975,0.975], ha='right', va='top',
                    xycoords='axes fraction', fontsize=11, bbox=bbox_props)

        # Mod on 20/11/2017
        fig.set_size_inches(8,6)
        fig.savefig(pp, format='pdf', bbox_inches='tight')

        # + on 20/11/2017, Mod on 18/03/2018
        mylog.info('Writing : '+out_fwhm_pdf)
        pp.close()

        # Compute transparency and plot | + on 09/12/2017
        out_trans_pdf = rawdir+prefix+'_stack_trans.pdf'

        pp = PdfPages(out_trans_pdf)

        trans0 = np.zeros(n_files)
        spec0  = np.zeros((n_files, naxis2))

        fig, ax = plt.subplots()

        lstyle0 = ['solid','dashed','dashdot','dotted',
                   'solid','dashed','dashdot','dotted']
        lw0 = [0.25, 0.33, 0.50, 0.75, 0.25, 0.33, 0.50, 0.75]

        for ii in range(n_files):
            x0    = np.arange(naxis1)
            t_sig = FWHM0[ii] / (2*np.sqrt(2*np.log(2))) / pscale
            idx   = np.where(np.abs(x0 - peak_val[ii])/t_sig <= 2.5)[0]

            im0        = diff_cube0[ii] / exptime0[ii]
            spec0[ii]  = np.sum(im0[:,idx], axis=1)
            trans0[ii] = np.sum(im0[:,idx])

            lstyle, lw = lstyle0[ii / 7], lw0[ii / 7]
            ax.plot(x0, spec0[ii], linewidth=lw, label=seqno[ii], alpha=0.5,
                    linestyle=lstyle)

        ax.set_xlim([0,2100])
        ax.set_xlabel('X [pixels]')
        ax.set_ylabel('Flux [ADU/s]')
        ax.legend(loc='upper left', fontsize='9', ncol=3, framealpha=0.5,
                  columnspacing=0.5)

        fig.set_size_inches(8,6)
        fig.savefig(pp, format='pdf', bbox_inches='tight')

        fig, ax = plt.subplots() # + on 20/11/2017
        i_max = np.argmax(trans0)
        ax.plot(seqno, trans0/trans0[i_max], marker='o', color='b', alpha=0.5)
        ax.set_xlabel('Image Frame No.')
        ax.set_ylabel('Relative Transparency')
        ax.set_ylim([0,1.1])
        ax.annotate(label0, [0.025,0.975], ha='left', va='top', weight='bold',
                    xycoords='axes fraction', fontsize=11, bbox=bbox_props)
        ax.minorticks_on()

        # Mod on 20/11/2017
        fig.set_size_inches(8,6)
        fig.savefig(pp, format='pdf', bbox_inches='tight')

        # + on 20/11/2017, Mod on 18/03/2018
        mylog.info('Writing : '+out_trans_pdf)
        pp.close()

        # + on 18/12/2018
        fwhm_wht0 = 1/FWHM0**2
        stack0_wht = np.ma.average(shift_cube0_mask, axis=0, weights=fwhm_wht0)

        out_fits_wht = rawdir+prefix+'_stack.wht.fits'
        mylog.info('Writing : '+out_fits_wht)
        fits.writeto(out_fits_wht, stack0_wht.data, overwrite=True)

    mylog.info('End main ! ') # Mod on 18/03/2018
#enddef

max_stars_repo_name: astrochun/MMTTools

max_issues_repo_name: astrochun/MMTTools

max_forks_repo_name: astrochun/MMTTools

reason: import numpy,from scipy,import astropy,from astropy


text:
"""Main Filter class."""
import enum

from dataclasses import dataclass, field
from typing import Iterable, NamedTuple

import numpy as np
import xarray as xr

from scipy import interpolate

from .gpu_compat import get_array_module
from .kernels import ALL_KERNELS, BaseLaplacian, GridType

FilterShape = enum.Enum("FilterShape", ["GAUSSIAN", "TAPER"])

class TargetSpec(NamedTuple):
    s_max: float
    filter_scale: float
    transition_width: float

# these functions return functions
def _gaussian_target(target_spec: TargetSpec):
    return lambda t: np.exp(
        -(target_spec.s_max * (t + 1) / 2) * (target_spec.filter_scale) ** 2 / 24
    )

def _taper_target(target_spec: TargetSpec):
    FK = interpolate.PchipInterpolator(
        np.array(
            [
                0,
                2 * np.pi / (target_spec.transition_width * target_spec.filter_scale),
                2 * np.pi / target_spec.filter_scale,
                2 * np.sqrt(target_spec.s_max),
            ]
        ),
        np.array([1, 1, 0, 0]),
    )
    return lambda t: FK(np.sqrt((t + 1) * (target_spec.s_max / 2)))

_target_function = {
    FilterShape.GAUSSIAN: _gaussian_target,
    FilterShape.TAPER: _taper_target,
}

class FilterSpec(NamedTuple):
    n_lap_steps: int
    s_l: Iterable[float]
    n_bih_steps: int
    s_b: Iterable[complex]

def _compute_filter_spec(
    filter_scale,
    dx_min,
    filter_shape,
    transition_width,
    ndim,
    n_steps=0,
    root_tolerance=1e-12,
):
    # First set number of steps if not supplied by user
    if n_steps == 0:
        if ndim > 2:
            raise ValueError(f"When ndim > 2, you must set n_steps manually")
        if filter_shape == FilterShape.GAUSSIAN:
            if ndim == 1:
                n_steps = np.ceil(1.3 * filter_scale / dx_min).astype(int)
            else:  # ndim==2
                n_steps = np.ceil(1.8 * filter_scale / dx_min).astype(int)
        else:  # Taper
            if ndim == 1:
                n_steps = np.ceil(4.5 * filter_scale / dx_min).astype(int)
            else:  # ndim==2
                n_steps = np.ceil(6.4 * filter_scale / dx_min).astype(int)

    # First set up the mass matrix for the Galerkin basis from Shen (SISC95)
    M = (np.pi / 2) * (
        2 * np.eye(n_steps - 1)
        - np.diag(np.ones(n_steps - 3), 2)
        - np.diag(np.ones(n_steps - 3), -2)
    )
    M[0, 0] = 3 * np.pi / 2

    # The range of wavenumbers is 0<=|k|<=sqrt(ndim)*pi/dxMin.
    # Per the notes, define s=k^2.
    # Need to rescale to t in [-1,1]: t = (2/sMax)*s -1; s = sMax*(t+1)/2
    s_max = ndim * (np.pi / dx_min) ** 2

    target_spec = TargetSpec(s_max, filter_scale, transition_width)
    F = _target_function[filter_shape](target_spec)

    # Compute inner products of Galerkin basis with target
    b = np.zeros(n_steps - 1)
    points, weights = np.polynomial.chebyshev.chebgauss(n_steps + 1)
    for i in range(n_steps - 1):
        tmp = np.zeros(n_steps + 1)
        tmp[i] = 1
        tmp[i + 2] = -1
        phi = np.polynomial.chebyshev.chebval(points, tmp)
        b[i] = np.sum(
            weights * phi * (F(points) - ((1 - points) / 2 + F(1) * (points + 1) / 2))
        )

    # Get polynomial coefficients in Galerkin basis
    c_hat = np.linalg.solve(M, b)
    # Convert back to Chebyshev basis coefficients
    p = np.zeros(n_steps + 1)
    p[0] = c_hat[0] + (1 + F(1)) / 2
    p[1] = c_hat[1] - (1 - F(1)) / 2
    for i in range(2, n_steps - 1):
        p[i] = c_hat[i] - c_hat[i - 2]
    p[n_steps - 1] = -c_hat[n_steps - 3]
    p[n_steps] = -c_hat[n_steps - 2]

    # Get roots of the polynomial
    r = np.polynomial.chebyshev.chebroots(p)

    # convert back to s in [0,sMax]
    s = s_max / 2 * (r + 1)
    # Separate out the real and complex roots
    n_lap_steps = np.size(s[np.where(np.abs(np.imag(r)) < root_tolerance)])
    s_l = np.real(s[np.where(np.abs(np.imag(r)) < root_tolerance)])
    n_bih_steps = (n_steps - n_lap_steps) // 2
    s_b_re, indices = np.unique(
        np.real(s[np.where(np.abs(np.imag(r)) > root_tolerance)]), return_index=True
    )
    s_b_im = np.imag(s[np.where(np.abs(np.imag(r)) > root_tolerance)])[indices]
    s_b = s_b_re + s_b_im * 1j

    return FilterSpec(n_lap_steps, s_l, n_bih_steps, s_b)

def _create_filter_func(
    filter_spec: FilterSpec,
    Laplacian: BaseLaplacian,
):
    """Returns a function whose first argument is the field to be filtered
    and whose subsequent arguments are the require grid variables
    """

    def filter_func(field, *args):
        # these next steps are a kind of hack we have to turn keyword arugments into regular arguments
        # the reason for doing this is that Xarray's apply_ufunc machinery works a lot better
        # with regular arguments
        assert len(args) == len(Laplacian.required_grid_args())
        grid_vars = {k: v for k, v in zip(Laplacian.required_grid_args(), args)}
        laplacian = Laplacian(**grid_vars)
        np = get_array_module(field)
        field_bar = field.copy()  # Initalize the filtering process
        for i in range(filter_spec.n_lap_steps):
            s_l = filter_spec.s_l[i]
            tendency = laplacian(field_bar)  # Compute Laplacian
            field_bar += (1 / s_l) * tendency  # Update filtered field
        for i in range(filter_spec.n_bih_steps):
            s_b = filter_spec.s_b[i]
            temp_l = laplacian(field_bar)  # Compute Laplacian
            temp_b = laplacian(temp_l)  # Compute Biharmonic (apply Laplacian twice)
            field_bar += (
                temp_l * 2 * np.real(s_b) / np.abs(s_b) ** 2
                + temp_b * 1 / np.abs(s_b) ** 2
            )
        return field_bar

    return filter_func

@dataclass
class Filter:
    """A class for applying diffusion-based smoothing filters to gridded data.

    ̦Parameters
    ----------
    filter_scale : float
        The filter scale, which has different meaning depending on filter shape
    dx_min : float
        The smallest grid spacing. Should have same units as ``filter_scale``
    n_steps : int, optional
        Number of total steps in the filter
        ``n_steps == 0`` means the number of steps is chosen automatically
    filter_shape : FilterShape
        - ``FilterShape.GAUSSIAN``: The target filter has kernel :math:`e^{-|x/Lf|^2}`
        - ``FilterShape.TAPER``: The target filter has target grid scale Lf. Smaller scales are zeroed out.
          Scales larger than ``pi * filter_scale / 2`` are left as-is. In between is a smooth transition.
    transition_width : float, optional
        Width of the transition region in the "Taper" filter.
    ndim : int, optional
         Laplacian is applied on a grid of dimension ndim
    grid_type : GridType
        what sort of grid we are dealing with
    grid_vars : dict
        dictionary of extra parameters used to initialize the grid laplacian

    Attributes
    ----------
    filter_spec: FilterSpec
    """

    filter_scale: float
    dx_min: float
    filter_shape: FilterShape = FilterShape.GAUSSIAN
    transition_width: float = np.pi
    ndim: int = 2
    n_steps: int = 0
    grid_type: GridType = GridType.CARTESIAN
    grid_vars: dict = field(default_factory=dict, repr=False)

    def __post_init__(self):

        if self.n_steps < 0:
            raise ValueError("Filter requires N>=0")

        self.filter_spec = _compute_filter_spec(
            self.filter_scale,
            self.dx_min,
            self.filter_shape,
            self.transition_width,
            self.ndim,
            self.n_steps,
        )

        # check that we have all the required grid aguments
        self.Laplacian = ALL_KERNELS[self.grid_type]
        if not set(self.Laplacian.required_grid_args()) == set(self.grid_vars):
            raise ValueError(
                f"Provided `grid_vars` {list(self.grid_vars)} do not match expected "
                f"{list(self.Laplacian.required_grid_args())}"
            )
        self.grid_ds = xr.Dataset({name: da for name, da in self.grid_vars.items()})

    def apply(self, field, dims):
        """Filter a field across the dimensions specified by dims."""

        filter_func = _create_filter_func(self.filter_spec, self.Laplacian)
        grid_args = [self.grid_ds[name] for name in self.Laplacian.required_grid_args()]
        assert len(dims) == 2
        n_args = 1 + len(grid_args)
        field_smooth = xr.apply_ufunc(
            filter_func,
            field,
            *grid_args,
            input_core_dims=n_args * [dims],
            output_core_dims=[dims],
            output_dtypes=[field.dtype],
            dask="parallelized",
        )
        return field_smooth

max_stars_repo_name: sdbachman/gcm-filters

max_issues_repo_name: sdbachman/gcm-filters

max_forks_repo_name: sdbachman/gcm-filters

reason: import numpy,from scipy


text:
# EXAMPLE METRICS CALCULATIONS

import numpy as np
import pandas as pd

#RESULTS DATA FORMAT

# contains one row for each simulated event, identifies if the line broke
breaks_df = pd.read_csv('sample_breaks_results.csv', index_col = 0)

# contains one row for each simulated event, and results for each municipality
sll_df = pd.read_csv('sample_sll_results.csv', index_col = 0)
vwl_df = pd.read_csv('sample_vwl_results.csv', index_col = 0)
vewl_df = pd.read_csv('sample_vewl_results.csv', index_col = 0)

# calculate metrics
line_sll_raw = np.dot(breaks_df.transpose(),sll_df.sum(axis= 1))/breaks_df.sum()
line_vwl_raw = np.dot(breaks_df.transpose(),vwl_df.sum(axis= 1))/breaks_df.sum()
line_vewl_raw = np.dot(breaks_df.transpose(),vewl_df.sum(axis= 1))/breaks_df.sum()

line_sll = line_sll_raw.rank()
line_vwl = line_vwl_raw.rank()
line_vewl = line_vewl_raw.rank()

max_stars_repo_path: CBES Stochastic Power Network Failure/Example Metrics Calculation/MetricsCalculation.py


text:
# -*- coding: utf-8 -*-

"""
Created on Mon, May 3, 2021
Almost exact MC for Heston model
@author: Shang Chencheng & Ning Lei
"""

import numpy as np

from scipy.special import iv
from scipy.misc import derivative
from scipy.stats import invgauss

class HestonMCAe:
    """
    Almost exact MC for Heston model.

    Underlying price is assumed to follow a geometric Brownian motion.
    Volatility (variance) of the price is assumed to follow a CIR process.
    Example:
        >>> strike = 100
        >>> spot = 100
        >>> vov, kappa, rho, texp, theta, sigma = 0.61, 6.21, -0.7, 1, 0.019, 0.10201
        >>> heston_ae = HestonMCAe(vov, kappa, rho, theta, r)
        >>> price_ae = heston_ae.price(strike, spot, texp, sigma_0, intr=0, divr=0)
        >>> price_ae
        8.946951375550809
    """
    def __init__(self, vov=1, kappa=0.5, rho=-0.9, theta=0.04, r=0):
        """
        Initiate a Heston model

        Args:
            vov: volatility of variance, strictly positive
            kappa: speed of variance's mean-reversion, strictly positive
            rho: correlation between BMs of price and vol
            theta: long-term mean (equilibirum level) of the variance, strictly positive
            r：the drift item
        """
        self.vov = vov
        self.theta = theta
        self.kappa = kappa
        self.rho = rho
        self.r = r

    def ch_f(self, texp, sigma_0, sigma_t, chi_dim):
        """

        Args:
            texp: time to expiry
            sigma_0: initial volatility
            sigma_t: volatility at time T
            chi_dim: dimensions of chisquare distribution

        Returns:
            ch_f: characteristic function of the distribution of integral sigma_t

        """
        gamma_f = lambda a: np.sqrt(self.kappa ** 2 - 2 * self.vov**2 * a * 1j)

        temp_f = lambda a: gamma_f(a) * texp

        ch_f_part_1 = lambda a: gamma_f(a) * np.exp(-0.5 * (temp_f(a) - self.kappa * texp)) \
                                * (1 - np.exp(-self.kappa * texp)) / (self.kappa * (1 - np.exp(-temp_f(a))))

        ch_f_part_2 = lambda a: np.exp((sigma_0**2 + sigma_t**2) / self.vov ** 2 * \
                                       (self.kappa * (1 + np.exp(-self.kappa * texp)) / (1 - np.exp(-self.kappa * texp))
                                        - gamma_f(a) * (1 + np.exp(-temp_f(a))) / (1 - np.exp(-temp_f(a)))))

        ch_f_part_3 = lambda a: iv(0.5 * chi_dim - 1, np.sqrt(sigma_0**2 * sigma_t**2) * 4 * gamma_f(a) *
                                   np.exp(-0.5 * temp_f(a)) / (self.vov ** 2 * (1 - np.exp(-temp_f(a))))) / \
                                iv(0.5 * chi_dim - 1, np.sqrt(sigma_0**2 * sigma_t**2) * 4 * self.kappa *
                                   np.exp(-0.5 * self.kappa * texp) / (
                                           self.vov ** 2 * (1 - np.exp(- self.kappa * texp))))

        ch_f = lambda a: ch_f_part_1(a) * ch_f_part_2(a) * ch_f_part_3(a)
        return ch_f

    def gen_vov_t(self, chi_dim, chi_lambda, texp, n_paths):
        """

        Args:
            chi_dim: dimensions of chisquare distribution
            chi_lambda: the skewing item of chisquare distribution
            texp: time to expiry
            n_paths: number of vol paths generated

        Returns:
            sigma_t: volatility at time T

        """
        cof = self.vov ** 2 * (1 - np.exp(-self.kappa * texp)) / (4 * self.kappa)
        sigma_t = np.sqrt(cof * np.random.noncentral_chisquare(chi_dim, chi_lambda, n_paths))
        return sigma_t

    def gen_s_t(self, spot, sigma_t, sigma_0, texp, integral_sigma_t, n_paths):
        """

        Args:
            spot: spot (or forward)
            sigma_t: volatility at time T
            sigma_0: initial volatility
            texp: time to expiry
            integral_sigma_t: samples from the distribution of integral sigma_t
            n_paths: number of vol paths generated

        Returns:
            s_t: stock price at time T
        """

        integral_sqrt_sigma_t = (sigma_t**2 - sigma_0**2 - self.kappa * self.theta * texp + self.kappa * integral_sigma_t)\
                                / self.vov
        mean = np.log(spot) + (self.r * texp - 0.5 * integral_sigma_t + self.rho * integral_sqrt_sigma_t)
        sigma_2 = np.sqrt((1 - self.rho ** 2) * integral_sigma_t)
        s_t = np.exp(mean + sigma_2 * np.random.normal(size=n_paths))
        return s_t

    def price(self, strike, spot, texp, sigma_0, intr=0, divr=0, n_paths=10000, seed=None,
              dis_can="Inverse-Gaussian", call=1):
        """
        Args:
            strike: strike price
            spot: spot (or forward)
            texp: time to expiry
            sigma_0: initial volatility
            intr: interest rate (domestic interest rate)
            divr: dividend/convenience yield (foreign interest rate)
            n_paths: number of vol paths generated
            seed: random seed for rv generation
        Returns:
            price_ae: option price
        """
        if seed:
            np.random.seed(seed)

        chi_dim = (4 * self.theta * self.kappa) / (self.vov ** 2)
        chi_lambda = (4 * self.kappa * np.exp(-self.kappa * texp)) / \
                     ((self.vov ** 2) * (1 - np.exp(-self.kappa * texp))) * sigma_0**2

        sigma_t = self.gen_vov_t(chi_dim, chi_lambda, texp, n_paths)

        ch_f = self.ch_f(texp, sigma_0, sigma_t, chi_dim)

        moment_1st = (derivative(ch_f, 0, n=1, dx=1e-5) / 1j).real
        moment_2st = (derivative(ch_f, 0, n=2, dx=1e-5) / (1j ** 2)).real

        if dis_can == "Inverse-Gaussian":
            scale_ig = moment_1st**3 / (moment_2st - moment_1st**2)
            miu_ig = moment_1st / scale_ig
            integral_sigma_t = invgauss.rvs(miu_ig, scale=scale_ig)
            s_t = self.gen_s_t(spot, sigma_t, sigma_0, texp, integral_sigma_t, n_paths)

        elif dis_can == "Log-normal":
            scale_ln = np.sqrt(np.log(moment_2st) - 2 * np.log(moment_1st))
            miu_ln = np.log(moment_1st) - 0.5 * scale_ln ** 2
            integral_sigma_t = np.random.lognormal(miu_ln, scale_ln)
            s_t = self.gen_s_t(spot, sigma_t, sigma_0, texp, integral_sigma_t, n_paths)
        else:
            print("This function is not currently a candidate function!")
            return -1

        if call:
            price_ae = np.fmax(s_t - strike, 0).mean()
        else:
            price_ae = np.fmax(strike - s_t, 0).mean()

        return np.exp(- self.r * texp) * price_ae

max_stars_repo_name: XueyangHu/ASP

max_issues_repo_name: XueyangHu/ASP

max_forks_repo_name: XueyangHu/ASP


text:
import random
import numbers
import collections
import math
import cv2
import numpy as np

from .crop import RandomCrop

from PIL import Image

class FreeScale(object):

  def __init__(self, size):
    self.size = tuple(reversed(size))  # size: (h, w)

  def __call__(self, img, mask):
    assert img.size == mask.size
    return (img.resize(self.size,
                       Image.BILINEAR), mask.resize(self.size, Image.NEAREST))

class RandomScale(object):

  def __init__(self, scale, aspect_ratio=None):
    assert (isinstance(scale, collections.Iterable) and len(scale) == 2)
    if isinstance(scale, collections.Iterable) and len(scale) == 2 \
        and isinstance(scale[0], numbers.Number) and isinstance(scale[1], numbers.Number) \
        and 0 < scale[0] < scale[1]:
      self.scale = scale
    else:
      raise (RuntimeError("RandomScale() scale param error.\n"))
    if aspect_ratio is None:
      self.aspect_ratio = aspect_ratio
    elif isinstance(aspect_ratio, collections.Iterable) and len(aspect_ratio) == 2 \
        and isinstance(aspect_ratio[0], numbers.Number) and isinstance(aspect_ratio[1], numbers.Number) \
        and 0 < aspect_ratio[0] < aspect_ratio[1]:
      self.aspect_ratio = aspect_ratio
    else:
      raise (RuntimeError("RandomScale() aspect_ratio param error.\n"))

  def __call__(self, image, label):
    image = np.asarray(image)
    label = np.asarray(label)
    temp_scale = self.scale[0] + (self.scale[1] -
                                  self.scale[0]) * random.random()
    temp_aspect_ratio = 1.0
    if self.aspect_ratio is not None:
      temp_aspect_ratio = self.aspect_ratio[0] + (
          self.aspect_ratio[1] - self.aspect_ratio[0]) * random.random()
      temp_aspect_ratio = math.sqrt(temp_aspect_ratio)
    scale_factor_x = temp_scale * temp_aspect_ratio
    scale_factor_y = temp_scale / temp_aspect_ratio
    image = cv2.resize(
        image,
        None,
        fx=scale_factor_x,
        fy=scale_factor_y,
        interpolation=cv2.INTER_LINEAR)
    label = cv2.resize(
        label,
        None,
        fx=scale_factor_x,
        fy=scale_factor_y,
        interpolation=cv2.INTER_NEAREST)
    return image, label

class RandomSized(object):

  def __init__(self, size):
    self.size = size
    self.scale = RandomScale(self.size)
    self.crop = RandomCrop(self.size)

  def __call__(self, img, mask):
    assert img.size == mask.size

    w = int(random.uniform(0.5, 2) * img.size[0])
    h = int(random.uniform(0.5, 2) * img.size[1])

    img, mask = (img.resize((w, h),
                            Image.BILINEAR), mask.resize((w, h), Image.NEAREST))

    return self.crop(*self.scale(img, mask))

class Resize(object):

  def __init__(self, size):
    assert (isinstance(size, collections.Iterable) and len(size) == 2)
    self.size = size

  def __call__(self, image, label):
    image = cv2.resize(image, self.size[::-1], interpolation=cv2.INTER_LINEAR)
    label = cv2.resize(label, self.size[::-1], interpolation=cv2.INTER_NEAREST)
    return image, label

max_stars_repo_name: VincentXWD/codebase-segmentation

max_issues_repo_name: VincentXWD/codebase-segmentation

max_forks_repo_name: VincentXWD/codebase-segmentation


text:
#
# Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <angelos.katharopoulos@idiap.ch>,
# Apoorv Vyas <avyas@idiap.ch>
#

import os
import time
import unittest

import numpy as np
import torch

from fast_transformers.causal_product.causal_product_cpu import \
    causal_dot_product, causal_dot_backward

def max_relative_error(a, b, eps=1e-6):
    return torch.abs((a-b) / (torch.abs(a) + eps)).max().item()

class TestCausalProductCPU(unittest.TestCase):
    kernels = {
        "normal": {
            "forward": causal_dot_product,
            "backward": causal_dot_backward
        },
        # map the optimized versions to other keys here
    }

    def _test_result_forward(self, CP):
        for t in range(10):
            N = 10
            L = 100
            H = 10
            E = np.random.randint(10, 256)
            M = np.random.randint(10, 256)
            Q = torch.rand(N, H, L, E)
            K = torch.rand(N, H, L, E)
            V = torch.rand(N, H, L, M)
            out = torch.zeros(N, H, L, M)
            self.kernels[CP]["forward"](Q, K, V, out)

            QK = torch.einsum("nhle,nhse->nhls", Q, K)
            mask = torch.tril(torch.ones(L, L))
            out2 = torch.einsum("nhls,nhsm,ls->nhlm", QK, V, mask)

            self.assertLess(max_relative_error(out2, out), 1e-5)

    def _test_result_backward(self, CP):
        for t in range(10):
            N = 10
            L = 100
            H = 10
            E = np.random.randint(10, 256)
            M = np.random.randint(10, 256)
            Q = torch.rand(N, H, L, E)
            K = torch.rand(N, H, L, E)
            V = torch.rand(N, H, L, M)
            go = torch.rand(N, H, L, M)
            gq, gk, gv = [torch.zeros_like(x) for x in [Q, K, V]]
            self.kernels[CP]["backward"](Q, K, V, go, gq, gk, gv)

            Q, K, V = [x.requires_grad_(True) for x in [Q, K, V]]
            QK = torch.einsum("nhle,nhse->nhls", Q, K)
            mask = torch.tril(torch.ones(L, L))
            out = torch.einsum("nhls,nhsm,ls->nhlm", QK, V, mask)
            l = torch.einsum("nhlm,nhlm->", out, go)
            l.backward()

            self.assertLess(max_relative_error(Q.grad, gq), 1e-5)
            self.assertLess(max_relative_error(K.grad, gk), 1e-5)
            self.assertLess(max_relative_error(V.grad, gv), 1e-5)

    def _test_benchmark_forward(self, CP):
        N = 10
        L = 1000
        H = 10
        E = 32
        M = 64
        Q = torch.rand(N, H, L, E)
        K = torch.rand(N, H, L, E)
        V = torch.rand(N, H, L, M)
        out = torch.rand(N, H, L, M)

        # warmup the cache
        for i in range(10):
            self.kernels[CP]["forward"](Q, K, V, out)

        # measure
        start = time.time()
        for i in range(10):
            self.kernels[CP]["forward"](Q, K, V, out)
        end = time.time()
        print("[{}] CPU time taken: {} (ms)".format(
            CP,
            (end-start)*1000
        ))

    def _test_benchmark_backward(self, CP):
        N = 10
        L = 1000
        H = 10
        E = 32
        M = 64
        Q = torch.rand(N, H, L, E)
        K = torch.rand(N, H, L, E)
        V = torch.rand(N, H, L, M)
        go = torch.rand(N, H, L, M)
        gq, gk, gv = [torch.zeros_like(x) for x in [Q, K, V]]

        # warmup the cache
        for i in range(10):
            self.kernels[CP]["backward"](Q, K, V, go, gq, gk, gv)

        # measure
        start = time.time()
        for i in range(10):
            self.kernels[CP]["backward"](Q, K, V, go, gq, gk, gv)
        end = time.time()
        print("[{}] CPU time taken: {} (ms)".format(
            CP,
            (end-start)*1000
        ))

    def test_result_forward(self):
        for k in self.kernels.keys():
            with self.subTest(msg=k):
                self._test_result_forward(k)

    def test_result_backward(self):
        for k in self.kernels.keys():
            with self.subTest(msg=k):
                self._test_result_backward(k)

    @unittest.skipUnless(os.getenv("BENCHMARK_TESTS", ""), "no benchmarks")
    def test_benchmark_forward(self):
        for k in self.kernels.keys():
            with self.subTest(msg=k):
                self._test_benchmark_forward(k)

    @unittest.skipUnless(os.getenv("BENCHMARK_TESTS", ""), "no benchmarks")
    def test_benchmark_backward(self):
        for k in self.kernels.keys():
            with self.subTest(msg=k):
                self._test_benchmark_backward(k)

if __name__ == "__main__":
    unittest.main()

max_stars_repo_name: SamuelCahyawijaya/fast-transformers

max_issues_repo_name: SamuelCahyawijaya/fast-transformers

max_forks_repo_name: SamuelCahyawijaya/fast-transformers


text:
import logging, matplotlib, os, sys
import anndata
import scanpy as sc
import numpy as np
import scipy as sp
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rcParams
from matplotlib import colors
import seaborn as sb
plt.rcParams['figure.figsize']=(8,8) #rescale figures
sc.settings.verbosity = 3
sc.set_figure_params(dpi=200, dpi_save=300)

adata = sc.read('raw_data.h5ad')
sc.pp.normalize_total(adata)
sc.pp.log1p(adata)
print(adata)

print('Number of cells: {:d}'.format(adata.n_obs))

sc.pp.highly_variable_genes(adata, flavor='cell_ranger', n_top_genes=2000)
sc.pl.highly_variable_genes(adata, show=False, save='highly_variable.pdf')

# Calculate the visualizations
sc.pp.pca(adata, n_comps=20, use_highly_variable=True, svd_solver='arpack') # PC=20 from Nature paper
sc.pp.neighbors(adata)
sc.tl.tsne(adata, n_jobs=3)
sc.tl.umap(adata, min_dist=0.6)
sc.tl.diffmap(adata)

sc.pl.pca_variance_ratio(adata, log=True, show=False, save='pca_variance.pdf')

# Perform clustering - using highly variable genes
sc.tl.leiden(adata, resolution=1.0, key_added='leiden_r1')
sc.tl.leiden(adata, resolution=0.5, key_added='leiden_r0.5')
sc.tl.leiden(adata, resolution=0.4, key_added='leiden_r0.4')
sc.tl.leiden(adata, resolution=0.35, key_added='leiden_r0.35')
sc.tl.leiden(adata, resolution=0.3, key_added='leiden_r0.3')
sc.tl.leiden(adata, resolution=0.25, key_added='leiden_r0.25')
sc.tl.leiden(adata, resolution=0.2, key_added='leiden_r0.2')
sc.tl.leiden(adata, resolution=0.1, key_added='leiden_r0.1')

adata.write('./learned.h5ad')

todraw = ['leiden_r1', 'leiden_r0.5', 'leiden_r0.4', 'leiden_r0.35', 'leiden_r0.3', 'leiden_r0.25', 'leiden_r0.2', 'leiden_r0.1', 'replicate']

#Visualize the clustering and how this is reflected by different technical covariates
sc.pl.tsne(adata, color=todraw, size=10, legend_loc='on data', show=False, save='tsne.pdf')
sc.pl.umap(adata, color=todraw, size=10, legend_loc='on data', show=False, save='umap.pdf')

reason: import numpy,import scipy


text:
import argparse
import os
from PIL import Image
import numpy as np
import copy
from glob import glob
import random
import cv2
import time
import multiprocessing as mp
from functools import partial
import darknet
import darknet_images

def check_arguments_errors(args):
    assert 0 < args.thresh < 1, "Threshold should be a float between zero and one (non-inclusive)"
    assert 0 < args.resolution, "Resolution should be an integer greater than zero"
    if not os.path.exists(args.input_folder):
        raise(ValueError("Invalid input folder path {}".format(os.path.abspath(args.input_folder))))
    if not os.path.exists(args.output_folder):
        raise(ValueError("Invalid output folder path {}".format(os.path.abspath(args.output_folder))))
    if not os.path.exists(args.config_file):
        raise(ValueError("Invalid config path {}".format(os.path.abspath(args.config_file))))
    if not os.path.exists(args.weights):
        raise(ValueError("Invalid weight path {}".format(os.path.abspath(args.weights))))
    if not os.path.exists(args.data_file):
        raise(ValueError("Invalid data file path {}".format(os.path.abspath(args.data_file))))

def check_batch_shape(images, batch_size):
    """
        Image sizes should be the same width and height
    """
    shapes = [image.shape for image in images]
    if len(set(shapes)) > 1:
        raise ValueError("Images don't have same shape")
    if len(shapes) > batch_size:
        raise ValueError("Batch size higher than number of images")
    return shapes[0]

def get_images(input_folder):

    images = []

    for ext in ('*.jpg', '*.jpeg', '*.png'):
        images.extend(glob(input_folder + ext))

    print('{} image files recognised in the provided input folder.'.format(len(images)))

    return images

def split_image(img_path, resolution):

    # Get image
    image = np.asarray(Image.open(img_path))
    w = image.shape[0]
    h = image.shape[1]

    # Store cropped images' names
    img_crop_names = []

    # Check if input resolution is high enough
    if w < args.resolution or h < args.resolution:
        print('Image |{}| resolution too low: {}x{}'.format(img_path, w, h))
        print('Please provide images with resolution at least {}x{}'.format(args.resolution, args.resolution))
        exit()

    # Determine number of sub-images to create
    nwi = int(np.ceil(w/args.resolution))
    wo = int(np.ceil((nwi*args.resolution - w)/(nwi - 1)))
    nhi = int(np.ceil(h/args.resolution))
    ho = int(np.ceil((nhi*args.resolution - h)/(nhi - 1)))

    # Generate crop images
    for i in range(nwi):
        for j in range(nhi):
            if i == nwi-1:
                i_s = w - args.resolution
                i_e = w
            else:
                i_s = i*args.resolution - i*wo
                i_e = i_s + args.resolution
            if j == nhi -1:
                j_s = h - args.resolution
                j_e = h
            else:
                j_s = j*args.resolution - j*ho
                j_e = j_s + args.resolution

            snip = image[i_s:i_e, j_s:j_e]
            Image.fromarray(snip).save(args.output_folder + os.path.splitext(os.path.basename(img_path))[0] + '_' + str(j_s) + '_' + str(i_s) + '.jpg')
            img_crop_names.append(args.output_folder + os.path.splitext(os.path.basename(img_path))[0] + '_' + str(j_s) + '_' + str(i_s) + '.jpg')

    return img_crop_names

def run_detection(output_splits, args):
    # Determine an optimal batch size
    if args.batch_size == None:
        args.batch_size = 1
        n_images = len(output_splits)
        while(n_images%(2*args.batch_size) == 0):
            args.batch_size *= 2
        n_batches = int(n_images/args.batch_size)

    print('Optimal batch size: {}'.format(args.batch_size))
    print('Loading network.')

    # Load network
    random.seed(3) # Deterministic bbox colours
    network, class_names, class_colors = darknet.load_network(
        args.config_file,
        args.data_file,
        args.weights,
        batch_size=args.batch_size
    )

    # Get network size
    n_w = darknet.network_width(network)
    n_h = darknet.network_height(network)

    # Split in batches
    global_predictions = []
    drawn_images = []
    for batch in range(n_batches):

        prev_time = time.time() # Get batch calculation start time

        # Get images in the batch
        names_batch = output_splits[batch*args.batch_size:((batch+1)*args.batch_size)]
        images_batch = [cv2.imread(n) for n in names_batch]
        images_batch = [cv2.resize(n, (n_w, n_h), interpolation=cv2.INTER_LINEAR) for n in images_batch] # Ensure image resolution is the same as network's
        h, w, _ = check_batch_shape(images_batch, args.batch_size)
        dnet_images = darknet_images.prepare_batch(images_batch, network)

        # Detect
        batch_detections = darknet.network_predict_batch(
            network,
            dnet_images,
            args.batch_size,
            w,
            h,
            args.thresh,
            args.hier_thresh,
            None,
            0,
            0
        )

        # Process detections
        batch_predictions = []
        for idx in range(args.batch_size):
            num = batch_detections[idx].num
            detections = batch_detections[idx].dets
            if args.nms:
                darknet.do_nms_obj(detections, num, len(class_names), args.nms)
            predictions = darknet.remove_negatives(detections, class_names, num)
            images_batch[idx] = darknet.draw_boxes(predictions, images_batch[idx], class_colors)
            batch_predictions.append(predictions)
        darknet.free_batch_detections(batch_detections, args.batch_size)

        # Store predictions and drawn images
        global_predictions.extend(batch_predictions)
        images_batch = [cv2.cvtColor(n, cv2.COLOR_BGR2RGB) for n in images_batch] # Ensure output is in RGB
        drawn_images.extend(images_batch)

        # Calculate FPS
        fps = int(args.batch_size/(time.time() - prev_time))
        print("FPS: {}".format(fps))

    return drawn_images, global_predictions, class_names

def main():

    # Get start time
    start_time = time.time()

    # Check argument errors
    check_arguments_errors(args)

    # Load data
    images = get_images(input_folder=args.input_folder)

    # Split images
    # TODO: Need to paralelize
    output_splits = []
    for i in images:

        # Split image
        output_splits.extend(split_image(i, args.resolution))

    print('{} output splits successfully genrated.'.format(len(output_splits)))
    print('Running detection.')

    # Run detection
    output_images, predictions, class_names = run_detection(output_splits, args)

    # Display and save results
    for img, pred, name in zip(output_images, predictions, output_splits):
        darknet_images.save_annotations(name, img, pred, class_names)
        # darknet.print_detections(pred, args.ext_output)

    # Output average processing time per image
    avg_proc_time = (time.time() - start_time)/len(images)
    print("Average processing time per image: {}".format(avg_proc_time))

if __name__ == "__main__":

    # TODO: Set defaults to automate servce/client
    # Define parser for input arguments
    parser = argparse.ArgumentParser(description='Split images in smaller ones')
    parser.add_argument('--input_folder', '-if', type=str)
    parser.add_argument('--output_folder', '-of', type=str)
    parser.add_argument('--resolution', '-res', type=int, default=720)
    parser.add_argument('--config_file', '-cfg', type=str)
    parser.add_argument('--data_file', '-df', type=str)
    parser.add_argument('--weights', '-w', type=str)
    parser.add_argument('--batch_size', '-bs', type=int)
    parser.add_argument('--thresh', '-t', type=float, default=0.25)
    parser.add_argument('--hier_thresh', '-ht', type=float, default=0.5)
    parser.add_argument('--nms', '-nms', type=float, default=0.45)
    parser.add_argument('--ext_output', '-ext', type=bool, default=True)
    args = parser.parse_args()

    main()

max_stars_repo_name: johnfaben/OCR

max_issues_repo_name: johnfaben/OCR

max_forks_repo_name: johnfaben/OCR


text:
# coding: utf-8

# ## Imputing missing values in sklearn
#
# Mean imputation replaces missing values with the mean value of that feature/variable. Mean imputation is one of the most 'naive' imputation methods because unlike more complex methods like k-nearest neighbors imputation, it does not use the information we have about an observation to estimate a value for it.

# In[1]:

import warnings
warnings.filterwarnings('ignore')

# In[2]:

import pandas as pd
import numpy as np
from sklearn.preprocessing import Imputer

# In[12]:

# Create an empty dataset
df = pd.DataFrame()

# Create two variables called x0 and x1. Make the first value of x1 a missing value
df['x0'] = [0.3051,0.4949,0.6974,np.nan,0.2231,np.nan,0.4436,0.5897,0.6308,0.5]
df['x1'] = [np.nan,0.2654,0.2615,0.5846,0.4615,np.nan,0.4962,0.3269,np.nan,0.6731]

# View the dataset
df

# In[18]:

# Create an imputer object that looks for 'Nan' values
mean_imputer = Imputer(missing_values='NaN', strategy='mean', axis=0)

# Train the imputor on the df dataset
mean_imputer = mean_imputer.fit(df)

# In[19]:

# Apply the imputer to the df dataset
imputed_df = mean_imputer.transform(df.values)

max_stars_repo_name: bjfisica/MachineLearning


text:
import itertools
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from seequery.embeddings.embeddings_manager import EmbeddingsManager
from seequery.ontology.ontology_manager import OntologyManager
from seequery.pipeline.match_item import MatchItem
from seequery.pipeline.pipeline_component import PipelineComponent
from seequery.pipeline.scored_translation import ScoredTranslation
from seequery.utils.meta_template import MetaTemplateChunks

class ContextualRescorer(PipelineComponent):
    """ Having possible translation, score them according to their coexistence. """
    def __init__(self, ontology_manager: OntologyManager, embeddings_mngr: EmbeddingsManager) -> None:
        self.ontology_mngr = ontology_manager
        self.embeddings_mngr = embeddings_mngr

    def process(self, data: dict) -> dict:
        """
            A method processing given data with current pipeline step.

            Args:
                data (dict): a dict holding object state

            Returns:
                dict: a dict holding object state enriched with current pipeline results
        """
        for meta_enriched_vocab in data['vocab_for_templates']:
            meta_enriched_vocab['swap'] = False

            if not meta_enriched_vocab['success']:
                continue

            if len(meta_enriched_vocab['meta'].relations) == 0:
                for chunk_idx, match_item in meta_enriched_vocab['vocab'].items():
                    match_item.scored_candidates = self.get_top_translations(match_item.scored_candidates)
            else:
                result, swap = self.get_best_combination(meta_enriched_vocab)
                pred_idx, lhs_idx, rhs_idx = self._get_required_idxs(meta_enriched_vocab['meta'])

                to_update = [(pred_idx, "best_property"), (lhs_idx, "best_lhs")]
                if rhs_idx:
                    to_update.append((rhs_idx, "best_rhs"))

                for idx, key in to_update:
                    if key not in result:
                        data['status'] = {
                            "type": "ERROR",
                            "message": f"No {key}"
                        }
                        return data
                    meta_enriched_vocab['vocab'][idx].scored_candidates = [result[key]['scored_translation']]
                if swap:
                    meta_enriched_vocab['swap'] = True
        return data

    def get_best_combination(self, meta_enriched_vocab: dict) -> Tuple[dict, Optional[bool]]:
        """ Get best combination.

            Args:
                meta_entriched_vocab (dict): extracted phrases and template metadata

            Returns:
                Tuple[dict, bool]: Best combination with argswitch info
        """
        best_property = None
        best_lhs = None
        best_rhs = None

        best_score = 0.0
        best_arg_swap = None

        for relation, lhs, rhs in self._construct_all_possible_connections(meta_enriched_vocab):
            # verify domain range restrictions
            relation_obj, relation_st = relation
            lhs_obj, lhs_st = lhs
            rhs_obj, rhs_st = rhs

            restriction_score, restriction_arg_swap = \
                self.ontology_mngr.calc_restriction_score(relation_obj, lhs_obj, rhs_obj)
            if restriction_score < 1.0:  # these candidates cannot be used together, skip
                continue

            scores = [relation_st.score, lhs_st.score]
            if rhs_st:
                scores.append(rhs_st.score)

            avg_translation_score = np.mean(scores)

            usage_score, usage_arg_swap = self.ontology_mngr.calc_usage_score(
                relation_obj, lhs_obj, rhs_obj)

            if restriction_arg_swap:
                arg_swap = restriction_arg_swap
            elif usage_arg_swap:
                arg_swap = usage_arg_swap
            else:
                arg_swap = False

            combined_score = 0.6 * avg_translation_score + 0.4 * usage_score

            if combined_score > best_score:
                best_score = combined_score
                best_arg_swap = arg_swap
                best_property = relation
                best_lhs = lhs
                best_rhs = rhs

        result = {}
        if best_property:
            result['best_property'] = {"obj": best_property[0], "scored_translation": best_property[1]}
        if best_lhs:
            result['best_lhs'] = {"obj": best_lhs[0], "scored_translation": best_lhs[1]}
        if best_rhs:
            result["best_rhs"] = {"obj": best_rhs[0], "scored_translation": best_rhs[1]}
        else:
            result["best_rhs"] = {"obj": None, "scored_translation": None}
        return result, best_arg_swap

    def _construct_all_possible_connections(self, meta_enriched_vocab: dict) -> Any:
        """ Construct all possible property and arguments translations.

            Args:
                meta_entriched_vocab (dict): extracted phrases and template metadata

            Returns:
                Any: An iterator over possible combinations
        """
        meta = meta_enriched_vocab['meta']
        vocab = meta_enriched_vocab['vocab']

        prop_idx, lhs_idx, rhs_idx = self._get_required_idxs(meta)

        property_translations = self._collect_ontology_objects(vocab, prop_idx)
        lhs_translations = self._collect_ontology_objects(vocab, lhs_idx)
        rhs_translations = self._collect_ontology_objects(vocab, rhs_idx) if rhs_idx else [(None, None)]

        return itertools.product(property_translations, lhs_translations, rhs_translations)

    def _get_required_idxs(self, meta: MetaTemplateChunks) -> Tuple[str, str, Optional[str]]:
        """ Get idxes for further processing.

            Args:
                meta (MetaTemplateChunks): meta information on template

            Returns:
                Tuple[str, str, Optional[str]]: List of chunk idxes
        """
        prop_idx = list(meta.relations)[0]
        entity_idxs = sorted(list(meta.entities))
        lhs_idx = entity_idxs[0]
        rhs_idx = entity_idxs[1] if len(entity_idxs) >= 2 else None
        return (prop_idx, lhs_idx, rhs_idx)

    def _collect_ontology_objects(self, vocab: Dict[str, MatchItem],
                                  chunk_idx: str) -> List[Tuple[Any, Optional[ScoredTranslation]]]:
        """ Get ontology objects from tranlations.

            Args:
                vocab (Dict[str, MatchItem]): chunk ids mapped to extracted vocab
                chunk_idx (str): idx of a chunk to be processed

            Returns:
                List[Tuple[Any, ScoredTranslation]]: List of objects assigned to a given chunk
        """
        objects: List[Tuple[Any, Optional[ScoredTranslation]]] = []
        for scored_translation in vocab[chunk_idx].scored_candidates:
            objects.append(
                (self.ontology_mngr.onto_map[scored_translation.category][scored_translation.onto_label],
                 scored_translation)
            )
        return objects

    def get_top_translations(self, scored_candidates: List[ScoredTranslation],
                             count: int = 1) -> Optional[List[ScoredTranslation]]:
        """ Get top /count/ scored translations.

            Args:
                scored candidates (List[ScoredTranslation]): a list of tranlations to extract from

            Returns:
                Optional[List[ScoredTranslation]]: top /count/ translations.
                                                   None if no translations available
        """
        if len(scored_candidates) > 0:
            return scored_candidates[:count]
        else:
            return None

max_stars_repo_name: SeeQuery/SeeQuery

max_issues_repo_name: SeeQuery/SeeQuery

max_forks_repo_name: SeeQuery/SeeQuery