text:
###############################################################################
# WaterTAP Copyright (c) 2021, The Regents of the University of California,
# through Lawrence Berkeley National Laboratory, Oak Ridge National
# Laboratory, National Renewable Energy Laboratory, and National Energy
# Technology Laboratory (subject to receipt of any required approvals from
# the U.S. Dept. of Energy). All rights reserved.
#
# Please see the files COPYRIGHT.md and LICENSE.md for full copyright and license
# information, respectively. These files are also available online at the URL
# "https://github.com/watertap-org/watertap/"
#
###############################################################################
import os
from pyomo.environ import (
    ConcreteModel,
    Block,
    Expression,
    value,
    TransformationFactory,
    units as pyunits,
    assert_optimal_termination,
)
from pyomo.network import Arc, SequentialDecomposition
from pyomo.util.check_units import assert_units_consistent

from idaes.core import FlowsheetBlock
from idaes.core.util import get_solver
from idaes.core.util.initialization import propagate_state

import idaes.core.util.scaling as iscale
from idaes.generic_models.unit_models import Mixer, Separator, Product, Feed
from idaes.generic_models.unit_models.mixer import MomentumMixingType
from idaes.generic_models.unit_models.translator import Translator
from idaes.generic_models.costing import UnitModelCostingBlock
from idaes.core.util.exceptions import ConfigurationError

from watertap.unit_models.pressure_exchanger import PressureExchanger
from watertap.unit_models.pressure_changer import Pump, EnergyRecoveryDevice
from watertap.core.util.initialization import assert_degrees_of_freedom

import watertap.property_models.seawater_prop_pack as prop_SW
from watertap.unit_models.reverse_osmosis_0D import (
    ReverseOsmosis0D,
    ConcentrationPolarizationType,
    MassTransferCoefficient,
    PressureChangeType,
)
from watertap.core.wt_database import Database
import watertap.core.zero_order_properties as prop_ZO
from watertap.unit_models.zero_order import (
    FeedZO,
    PumpElectricityZO,
    NanofiltrationZO,
    SecondaryTreatmentWWTPZO,
)
from watertap.core.zero_order_costing import ZeroOrderCosting
from watertap.costing import WaterTAPCosting

def main():
    m = build()
    set_operating_conditions(m)

    assert_degrees_of_freedom(m, 0)
    assert_units_consistent(m)

    initialize_system(m)
    assert_degrees_of_freedom(m, 0)

    results = solve(m)

    add_costing(m)
    initialize_costing(m)
    assert_degrees_of_freedom(m, 0)

    solve(m)
    display_results(m)
    display_costing(m)

    return m, results

def build():
    # flowsheet set up
    m = ConcreteModel()
    m.db = Database()

    m.fs = FlowsheetBlock(default={"dynamic": False})

    # define property packages
    m.fs.prop_nf = prop_ZO.WaterParameterBlock(default={"solute_list": ["dye", "tds"]})
    m.fs.prop_ro = prop_SW.SeawaterParameterBlock()

    # define blocks
    prtrt = m.fs.pretreatment = Block()
    dye_sep = m.fs.dye_separation = Block()
    desal = m.fs.desalination = Block()

    # define flowsheet inlets and outlets
    m.fs.feed = FeedZO(default={"property_package": m.fs.prop_nf})
    m.fs.wwt_retentate = Product(default={"property_package": m.fs.prop_nf})
    m.fs.dye_retentate = Product(default={"property_package": m.fs.prop_nf})
    m.fs.permeate = Product(default={"property_package": m.fs.prop_ro})
    m.fs.brine = Product(default={"property_package": m.fs.prop_ro})

    # pretreatment
    prtrt.wwtp = SecondaryTreatmentWWTPZO(
        default={
            "property_package": m.fs.prop_nf,
            "database": m.db,
            "process_subtype": "default",
        }
    )

    # nanofiltration components
    dye_sep.P1 = PumpElectricityZO(
        default={
            "property_package": m.fs.prop_nf,
            "database": m.db,
            "process_subtype": "default",
        }
    )

    dye_sep.nanofiltration = NanofiltrationZO(
        default={
            "property_package": m.fs.prop_nf,
            "database": m.db,
            "process_subtype": "rHGO_dye_rejection",
        }
    )

    # reverse osmosis components

    desal.P2 = Pump(default={"property_package": m.fs.prop_ro})
    desal.RO = ReverseOsmosis0D(
        default={
            "property_package": m.fs.prop_ro,
            "has_pressure_change": True,
            "pressure_change_type": PressureChangeType.calculated,
            "mass_transfer_coefficient": MassTransferCoefficient.calculated,
            "concentration_polarization_type": ConcentrationPolarizationType.calculated,
        }
    )

    desal.RO.width.setub(2000)
    desal.RO.area.setub(20000)

    desal.S1 = Separator(
        default={"property_package": m.fs.prop_ro, "outlet_list": ["P2", "PXR"]}
    )
    desal.M1 = Mixer(
        default={
            "property_package": m.fs.prop_ro,
            "momentum_mixing_type": MomentumMixingType.equality,  # booster pump will match pressure
            "inlet_list": ["P2", "P3"],
        }
    )
    desal.PXR = PressureExchanger(default={"property_package": m.fs.prop_ro})
    desal.P3 = Pump(default={"property_package": m.fs.prop_ro})

    # translator blocks
    m.fs.tb_nf_ro = Translator(
        default={
            "inlet_property_package": m.fs.prop_nf,
            "outlet_property_package": m.fs.prop_ro,
        }
    )

    # since the dye << tds: Assume RO_TDS = NF_tds + NF_dye
    @m.fs.tb_nf_ro.Constraint(["H2O", "dye"])
    def eq_flow_mass_comp(blk, j):
        if j == "dye":
            return (
                blk.properties_in[0].flow_mass_comp["dye"]
                + blk.properties_in[0].flow_mass_comp["tds"]
                == blk.properties_out[0].flow_mass_phase_comp["Liq", "TDS"]
            )
        else:
            return (
                blk.properties_in[0].flow_mass_comp["H2O"]
                == blk.properties_out[0].flow_mass_phase_comp["Liq", "H2O"]
            )

    # connections
    m.fs.s_feed = Arc(source=m.fs.feed.outlet, destination=prtrt.wwtp.inlet)
    prtrt.s01 = Arc(source=prtrt.wwtp.treated, destination=dye_sep.P1.inlet)
    prtrt.s02 = Arc(source=prtrt.wwtp.byproduct, destination=m.fs.wwt_retentate.inlet)
    dye_sep.s01 = Arc(
        source=dye_sep.P1.outlet, destination=dye_sep.nanofiltration.inlet
    )
    dye_sep.s02 = Arc(
        source=dye_sep.nanofiltration.byproduct, destination=m.fs.dye_retentate.inlet
    )
    m.fs.s_nf = Arc(
        source=dye_sep.nanofiltration.treated, destination=m.fs.tb_nf_ro.inlet
    )

    m.fs.s_ro = Arc(source=m.fs.tb_nf_ro.outlet, destination=desal.S1.inlet)
    desal.s01 = Arc(source=desal.S1.P2, destination=desal.P2.inlet)
    desal.s02 = Arc(source=desal.P2.outlet, destination=desal.M1.P2)
    desal.s03 = Arc(source=desal.M1.outlet, destination=desal.RO.inlet)
    desal.s04 = Arc(
        source=desal.RO.retentate, destination=desal.PXR.high_pressure_inlet
    )
    desal.s05 = Arc(source=desal.S1.PXR, destination=desal.PXR.low_pressure_inlet)
    desal.s06 = Arc(source=desal.PXR.low_pressure_outlet, destination=desal.P3.inlet)
    desal.s07 = Arc(source=desal.P3.outlet, destination=desal.M1.P3)
    m.fs.s_disposal = Arc(
        source=desal.PXR.high_pressure_outlet, destination=m.fs.brine.inlet
    )

    m.fs.s_permeate = Arc(source=desal.RO.permeate, destination=m.fs.permeate.inlet)

    TransformationFactory("network.expand_arcs").apply_to(m)

    # scaling
    m.fs.prop_ro.set_default_scaling("flow_mass_phase_comp", 1e-3, index=("Liq", "H2O"))
    m.fs.prop_ro.set_default_scaling("flow_mass_phase_comp", 1e-1, index=("Liq", "TDS"))

    # set unit model values
    iscale.set_scaling_factor(desal.P2.control_volume.work, 1e-5)
    iscale.set_scaling_factor(desal.RO.area, 1e-4)
    iscale.set_scaling_factor(desal.P3.control_volume.work, 1e-5)
    iscale.set_scaling_factor(desal.PXR.low_pressure_side.work, 1e-5)
    iscale.set_scaling_factor(desal.PXR.high_pressure_side.work, 1e-5)

    # calculate and propagate scaling factors
    iscale.calculate_scaling_factors(m)
    return m

def set_operating_conditions(m):
    prtrt = m.fs.pretreatment
    dye_sep = m.fs.dye_separation
    desal = m.fs.desalination

    # feed
    flow_vol = 120 / 3600 * pyunits.m**3 / pyunits.s
    conc_mass_dye = 2.5 * pyunits.kg / pyunits.m**3
    conc_mass_tds = 50.0 * pyunits.kg / pyunits.m**3
    temperature = 298 * pyunits.K
    pressure = 101325 * pyunits.Pa

    m.fs.feed.flow_vol[0].fix(flow_vol)
    m.fs.feed.conc_mass_comp[0, "dye"].fix(conc_mass_dye)
    m.fs.feed.conc_mass_comp[0, "tds"].fix(conc_mass_tds)
    solve(m.fs.feed)

    # pretreatment
    prtrt.wwtp.load_parameters_from_database(use_default_removal=True)

    # nanofiltration
    dye_sep.nanofiltration.load_parameters_from_database(use_default_removal=True)

    # nf pump
    dye_sep.P1.load_parameters_from_database(use_default_removal=True)
    dye_sep.P1.applied_pressure.fix(
        dye_sep.nanofiltration.applied_pressure.get_values()[0]
    )
    dye_sep.P1.lift_height.unfix()

    # desalination
    desal.P2.efficiency_pump.fix(0.80)
    operating_pressure = 70e5 * pyunits.Pa
    desal.P2.control_volume.properties_out[0].pressure.fix(operating_pressure)
    desal.RO.A_comp.fix(4.2e-12)  # membrane water permeability
    desal.RO.B_comp.fix(3.5e-8)  # membrane salt permeability
    desal.RO.channel_height.fix(1e-3)  # channel height in membrane stage [m]
    desal.RO.spacer_porosity.fix(0.97)  # spacer porosity in membrane stage [-]
    desal.RO.permeate.pressure[0].fix(pressure)  # atmospheric pressure [Pa]
    desal.RO.velocity[0, 0].fix(0.25)
    desal.RO.recovery_vol_phase[0, "Liq"].fix(0.5)
    m.fs.tb_nf_ro.properties_out[0].temperature.fix(temperature)
    m.fs.tb_nf_ro.properties_out[0].pressure.fix(pressure)

    # pressure exchanger
    desal.PXR.efficiency_pressure_exchanger.fix(0.95)
    # booster pump
    desal.P3.efficiency_pump.fix(0.80)

    return

def initialize_system(m):
    prtrt = m.fs.pretreatment
    dye_sep = m.fs.dye_separation
    desal = m.fs.desalination

    # initialize feed
    solve(m.fs.feed)

    # pretreatment
    propagate_state(m.fs.s_feed)
    s = SequentialDecomposition()
    s.options.tear_set = []
    s.options.iterLim = 1
    s.run(prtrt, lambda u: u.initialize())

    # initialized nf
    propagate_state(prtrt.s01)
    seq = SequentialDecomposition()
    seq.options.tear_set = []
    seq.options.iterLim = 1
    seq.run(dye_sep, lambda u: u.initialize())

    # initialize ro
    propagate_state(m.fs.s_nf)
    propagate_state(m.fs.s_ro)
    propagate_state(desal.s01)
    propagate_state(m.fs.s_disposal)
    propagate_state(m.fs.s_permeate)

    m.fs.tb_nf_ro.properties_out[0].flow_mass_phase_comp["Liq", "H2O"] = value(
        m.fs.tb_nf_ro.properties_in[0].flow_mass_comp["H2O"]
    )
    m.fs.tb_nf_ro.properties_out[0].flow_mass_phase_comp["Liq", "TDS"] = value(
        m.fs.tb_nf_ro.properties_in[0].flow_mass_comp["tds"]
    ) + value(m.fs.tb_nf_ro.properties_in[0].flow_mass_comp["dye"])

    desal.RO.feed_side.properties_in[0].flow_mass_phase_comp["Liq", "H2O"] = value(
        m.fs.tb_nf_ro.properties_out[0].flow_mass_phase_comp["Liq", "H2O"]
    )
    desal.RO.feed_side.properties_in[0].temperature = value(
        m.fs.tb_nf_ro.properties_out[0].temperature
    )
    desal.RO.feed_side.properties_in[0].pressure = value(
        desal.P2.control_volume.properties_out[0].pressure
    )
    solve(desal)
    desal.RO.initialize()
    return

def solve(blk, solver=None, tee=False, check_termination=True):
    if solver is None:
        solver = get_solver()
    results = solver.solve(blk, tee=tee)
    if check_termination:
        assert_optimal_termination(results)
    return results

def add_costing(m):
    prtrt = m.fs.pretreatment
    dye_sep = m.fs.dye_separation
    desal = m.fs.desalination

    # Zero order costing
    source_file = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "dye_desalination_global_costing.yaml",
    )

    m.fs.zo_costing = ZeroOrderCosting(default={"case_study_definition": source_file})
    m.fs.ro_costing = WaterTAPCosting()

    # cost nanofiltration module and pump
    prtrt.wwtp.costing = UnitModelCostingBlock(
        default={"flowsheet_costing_block": m.fs.zo_costing}
    )
    dye_sep.nanofiltration.costing = UnitModelCostingBlock(
        default={"flowsheet_costing_block": m.fs.zo_costing}
    )
    dye_sep.P1.costing = UnitModelCostingBlock(
        default={"flowsheet_costing_block": m.fs.zo_costing}
    )

    # RO Train
    # RO equipment is costed using more detailed costing package

    desal.P2.costing = UnitModelCostingBlock(
        default={
            "flowsheet_costing_block": m.fs.ro_costing,
            "costing_method_arguments": {"cost_electricity_flow": True},
        }
    )
    desal.RO.costing = UnitModelCostingBlock(
        default={"flowsheet_costing_block": m.fs.ro_costing}
    )

    desal.M1.costing = UnitModelCostingBlock(
        default={"flowsheet_costing_block": m.fs.ro_costing}
    )
    desal.PXR.costing = UnitModelCostingBlock(
        default={"flowsheet_costing_block": m.fs.ro_costing}
    )
    desal.P3.costing = UnitModelCostingBlock(
        default={
            "flowsheet_costing_block": m.fs.ro_costing,
            "costing_method_arguments": {"cost_electricity_flow": True},
        }
    )

    # Aggregate unit level costs and calculate overall process costs
    m.fs.zo_costing.cost_process()
    m.fs.ro_costing.cost_process()

    # Annual disposal of brine
    m.fs.brine_disposal_cost = Expression(
        expr=(
            m.fs.zo_costing.utilization_factor
            * (
                m.fs.zo_costing.waste_disposal_cost
                * pyunits.convert(
                    m.fs.brine.properties[0].flow_vol,
                    to_units=pyunits.m**3 / m.fs.zo_costing.base_period,
                )
            )
        ),
        doc="Cost of disposing of brine waste",
    )

    m.fs.sludge_disposal_cost = Expression(
        expr=(
            m.fs.zo_costing.utilization_factor
            * (
                m.fs.zo_costing.waste_disposal_cost
                * pyunits.convert(
                    m.fs.wwt_retentate.properties[0].flow_vol,
                    to_units=pyunits.m**3 / m.fs.zo_costing.base_period,
                )
            )
        ),
        doc="Cost of disposing of waste water treatment plant sludge",
    )

    m.fs.dye_recovery_revenue = Expression(
        expr=(
            m.fs.zo_costing.utilization_factor
            * m.fs.zo_costing.dye_mass_cost
            * pyunits.convert(
                m.fs.dye_retentate.flow_mass_comp[0, "dye"],
                to_units=pyunits.kg / m.fs.zo_costing.base_period,
            )
        ),
        doc="Savings from dye recovered back to the plant",
    )

    m.fs.water_recovery_revenue = Expression(
        expr=(
            m.fs.zo_costing.utilization_factor
            * m.fs.zo_costing.recovered_water_cost
            * pyunits.convert(
                m.fs.permeate.properties[0].flow_vol,
                to_units=pyunits.m**3 / m.fs.zo_costing.base_period,
            )
        ),
        doc="Savings from water recovered back to the plant",
    )

    # Combine results from costing packages and calculate overall metrics
    @m.Expression()
    def total_capital_cost(b, doc="Total capital cost of the treatment train"):
        return pyunits.convert(
            m.fs.zo_costing.total_capital_cost, to_units=pyunits.USD_2020
        ) + pyunits.convert(
            m.fs.ro_costing.total_investment_cost, to_units=pyunits.USD_2020
        )

    @m.Expression()
    def total_operating_cost(b, doc="Total operating cost of the treatment train"):
        return (
            pyunits.convert(
                m.fs.zo_costing.total_fixed_operating_cost,
                to_units=pyunits.USD_2020 / pyunits.year,
            )
            + pyunits.convert(
                m.fs.zo_costing.total_variable_operating_cost,
                to_units=pyunits.USD_2020 / pyunits.year,
            )
            + pyunits.convert(
                m.fs.ro_costing.total_operating_cost,
                to_units=pyunits.USD_2020 / pyunits.year,
            )
        )

    @m.Expression()
    def total_externalities(
        b, doc="Total cost of water/dye recovered and brine/sludge disposed"
    ):
        return pyunits.convert(
            m.fs.water_recovery_revenue
            + m.fs.dye_recovery_revenue
            - m.fs.brine_disposal_cost
            - m.fs.sludge_disposal_cost,
            to_units=pyunits.USD_2020 / pyunits.year,
        )

    @m.Expression()
    def LCOT(b, doc="Levelized cost of treatment with respect to volumetric feed flow"):
        return (
            b.total_capital_cost * b.fs.zo_costing.capital_recovery_factor
            + b.total_operating_cost
            - b.total_externalities
        ) / (
            pyunits.convert(
                b.fs.feed.properties[0].flow_vol,
                to_units=pyunits.m**3 / pyunits.year,
            )
            * b.fs.zo_costing.utilization_factor
        )

    @m.Expression()
    def LCOW(b, doc="Levelized cost of water with respect to volumetric permeate flow"):
        return (
            b.total_capital_cost * b.fs.zo_costing.capital_recovery_factor
            + b.total_operating_cost
            - pyunits.convert(
                m.fs.dye_recovery_revenue
                - m.fs.brine_disposal_cost
                - m.fs.sludge_disposal_cost,
                to_units=pyunits.USD_2020 / pyunits.year,
            )
        ) / (
            pyunits.convert(
                b.fs.permeate.properties[0].flow_vol,
                to_units=pyunits.m**3 / pyunits.year,
            )
            * b.fs.zo_costing.utilization_factor
        )

    assert_units_consistent(m)
    return

def initialize_costing(m):
    m.fs.zo_costing.initialize()
    m.fs.ro_costing.initialize()
    return

def display_results(m):
    print("Unit models:")
    m.fs.pretreatment.wwtp.report()
    m.fs.dye_separation.P1.report()
    m.fs.dye_separation.nanofiltration.report()
    m.fs.desalination.S1.report()
    m.fs.desalination.P2.report()
    m.fs.desalination.P3.report()
    m.fs.desalination.M1.report()
    m.fs.desalination.RO.report()
    m.fs.desalination.PXR.report()

    print("Streams:")
    flow_list = ["feed", "wwt_retentate", "dye_retentate", "permeate", "brine"]
    for f in flow_list:
        m.fs.component(f).report()
    return

def display_costing(m):
    print("\n Costing metrics:")
    capex = value(pyunits.convert(m.total_capital_cost, to_units=pyunits.MUSD_2020))

    opex = value(
        pyunits.convert(
            m.total_operating_cost, to_units=pyunits.MUSD_2020 / pyunits.year
        )
    )

    print("\n System costing metrics:")
    capex = value(pyunits.convert(m.total_capital_cost, to_units=pyunits.MUSD_2020))

    opex = value(
        pyunits.convert(
            m.total_operating_cost, to_units=pyunits.MUSD_2020 / pyunits.year
        )
    )

    externalities = value(
        pyunits.convert(
            m.total_externalities, to_units=pyunits.MUSD_2020 / pyunits.year
        )
    )

    lcot = value(pyunits.convert(m.LCOT, to_units=pyunits.USD_2020 / pyunits.m**3))

    lcow = value(pyunits.convert(m.LCOW, to_units=pyunits.USD_2020 / pyunits.m**3))

    print(f"Total Capital Cost: {capex:.4f} M$")

    print(f"Total Operating Cost: {opex:.4f} M$/year")

    print(f"Total Externalities: {externalities:.4f} M$/year")

    print(f"Levelized cost of treatment: {lcot:.4f} $/m3 feed")

    print(f"Levelized cost of water: {lcow:.4f} $/m3 permeate")

if __name__ == "__main__":
    model, results = main()

reason: from pyomo


text:
""" MIDI encoding base class and methods
TODO Control change messages (sustain, modulation, pitch bend)

"""

from sys import stdout
from pathlib import Path, PurePath
import json
from typing import List, Tuple, Dict, Union, Callable, Optional, Any

import numpy as np
from miditoolkit import MidiFile, Instrument, Note, TempoChange

from .constants import TIME_DIVISION, CHORD_MAPS

class Event:
    """ Event class, representing a token and its characteristics
    The name corresponds to the token type (e.g. Pitch, Position ...);
    The value to its value.
    These two attributes are used in events_to_tokens and tokens_to_events
    methods (see below) to convert an Event object to its corresponding integer.
    """

    def __init__(self, name, time, value, text):
        self.name = name
        self.time = time
        self.value = value
        self.text = text

    def __repr__(self):
        return f'Event(name={self.name}, time={self.time}, value={self.value}, text={self.text})'

class MIDITokenizer:
    """ MIDI encoding base class, containing common parameters to all encodings
    and common methods.

    :param pitch_range: range of used MIDI pitches
    :param beat_res: beat resolutions, with the form:
            {(beat_x1, beat_x2): beat_res_1, (beat_x2, beat_x3): beat_res_2, ...}
            The keys of the dict are tuples indicating a range of beats, ex 0 to 3 for the first bar
            The values are the resolution, in frames per beat, of the given range, ex 8
    :param nb_velocities: number of velocity bins
    :param additional_tokens: specifies additional tokens (chords, empty bars, tempo...)
    :param program_tokens: will add entries for MIDI programs in the dictionary, to use
            in the case of multitrack generation for instance
    :param params: can be a path to the parameter (json encoded) file or a dictionary
    """
    def __init__(self, pitch_range: range, beat_res: Dict[Tuple[int, int], int], nb_velocities: int,
                 additional_tokens: Dict[str, bool], program_tokens: bool,
                 params: Union[str, Path, PurePath, Dict[str, Any]] = None):
        if params is None:
            self.pitch_range = pitch_range
            self.beat_res = beat_res
            self.additional_tokens = additional_tokens
            self.nb_velocities = nb_velocities
        else:
            self.load_params(params)

        self.durations = self._create_durations_tuples()
        self.velocity_bins = np.linspace(0, 127, self.nb_velocities + 1, dtype=np.intc)
        np.delete(self.velocity_bins, 0)  # removes velocity 0
        if additional_tokens['Tempo']:
            self.tempo_bins = np.linspace(*additional_tokens['tempo_range'], additional_tokens['nb_tempos'],
                                          dtype=np.intc)
        else:
            self.tempo_bins = np.zeros(1)

        self.event2token, self.token2event, self.token_types_indices = self._create_vocabulary(program_tokens)

        # Keep in memory durations in ticks for seen time divisions so these values
        # are not calculated each time a MIDI is processed
        self.durations_ticks = {}

        # Holds the tempo changes, time signature, time division and key signature of a
        # MIDI (being parsed) so that methods processing tracks can access them
        self.current_midi_metadata = {}  # needs to be updated each time a MIDI is read

    def midi_to_tokens(self, midi: MidiFile) -> List[List[Union[int, List[int]]]]:
        """ Converts a MIDI file in a tokens representation.
        NOTE: if you override this method, be sure to keep every line of code below until
        the "Convert track to token" comment in the for loop

        :param midi: the MIDI objet to convert
        :return: the token representation, i.e. tracks converted into sequences of tokens
        """
        # Check if the durations values have been calculated before for this time division
        try:
            _ = self.durations_ticks[midi.ticks_per_beat]
        except KeyError:
            self.durations_ticks[midi.ticks_per_beat] = [(beat * res + pos) * midi.ticks_per_beat // res
                                                         for beat, pos, res in self.durations]

        # Register MIDI metadata
        self.current_midi_metadata = {'time_division': midi.ticks_per_beat,
                                      'tempo_changes': midi.tempo_changes,
                                      'time_sig_changes': midi.time_signature_changes,
                                      'key_sig_changes': midi.key_signature_changes}

        # Quantize tempo changes times
        quantize_tempos(midi.tempo_changes, midi.ticks_per_beat, max(self.beat_res.values()))

        tokens = []
        for track in midi.instruments:
            quantize_note_times(track.notes, self.current_midi_metadata['time_division'], max(self.beat_res.values()))
            track.notes.sort(key=lambda x: (x.start, x.pitch))  # sort notes
            remove_duplicated_notes(track.notes)  # remove possible duplicated notes

            # Convert track to tokens
            tokens.append(self.track_to_tokens(track))

        return tokens

    def track_to_tokens(self, track: Instrument) -> List[Union[int, List[int]]]:
        """ Converts a track (miditoolkit.Instrument object) into a sequence of tokens

        :param track: MIDI track to convert
        :return: sequence of corresponding tokens
        """
        raise NotImplementedError

    def events_to_tokens(self, events: List[Event]) -> List[int]:
        """ Converts a list of Event objects into a list of tokens
        You can override this method if necessary (e.g. CP Word)

        :param events: list of Events objects to convert
        :return: list of corresponding tokens
        """
        return [self.event2token[f'{event.name}_{event.value}'] for event in events]

    def tokens_to_events(self, tokens: List[int]) -> List[Event]:
        """ Convert a sequence of tokens in their respective event objects
        You can override this method if necessary (e.g. CP Word)

        :param tokens: sequence of tokens to convert
        :return: the sequence of corresponding events
        """
        events = []
        for token in tokens:
            name, val = self.token2event[token].split('_')
            events.append(Event(name, None, val, None))
        return events

    def tokens_to_midi(self, tokens: List[List[Union[int, List[int]]]],
                       programs: Optional[List[Tuple[int, bool]]] = None, output_path: Optional[str] = None,
                       time_division: Optional[int] = TIME_DIVISION) -> MidiFile:
        """ Convert multiple sequences of tokens into a multitrack MIDI and save it.
        The tokens will be converted to event objects and then to a miditoolkit.MidiFile object.
        NOTE: With Remi, MIDI-Like, CP Word or other encoding methods that process tracks
        independently, only the tempo changes of the first track in tokens will be used

        :param tokens: list of lists of tokens to convert, each list inside the
                       first list corresponds to a track
        :param programs: programs of the tracks
        :param output_path: path to save the file (with its name, e.g. music.mid),
                        leave None to not save the file
        :param time_division: MIDI time division / resolution, in ticks/beat (of the MIDI to create)
        :return: the midi object (miditoolkit.MidiFile)
        """
        midi = MidiFile(ticks_per_beat=time_division)
        for i, track_tokens in enumerate(tokens):
            if programs is not None:
                track, tempo_changes = self.tokens_to_track(track_tokens, time_division, programs[i])
            else:
                track, tempo_changes = self.tokens_to_track(track_tokens, time_division)
            midi.instruments.append(track)
            if i == 0:  # only keep tempo changes of the first track
                midi.tempo_changes = tempo_changes
                midi.tempo_changes[0].time = 0

        # Write MIDI file
        if output_path:
            Path(output_path).mkdir(parents=True, exist_ok=True)
            midi.dump(output_path)
        return midi

    def tokens_to_track(self, tokens: List[Union[int, List[int]]], time_division: Optional[int] = TIME_DIVISION,
                        program: Optional[Tuple[int, bool]] = (0, False)) -> Tuple[Instrument, List[TempoChange]]:
        """ Converts a sequence of tokens into a track object

        :param tokens: sequence of tokens to convert
        :param time_division: MIDI time division / resolution, in ticks/beat (of the MIDI to create)
        :param program: the MIDI program of the produced track and if it drum, (default (0, False), piano)
        :return: the miditoolkit instrument object and the possible tempo changes
        """
        raise NotImplementedError

    def _create_vocabulary(self, program_tokens: bool) -> Tuple[dict, dict, dict]:
        """ Create the tokens <-> event dictionaries
        These dictionaries are created arbitrary according to constants defined
        at the top of this file.
        Note that when using them (prepare_data method), there is no error-handling
        so you must be sure that every case is covered by the dictionaries.
        NOTE: token index 0 is often used as a padding index during training, it might
        be preferable to leave it as it to pad your batch sequences

        :param program_tokens: creates tokens for MIDI programs in the dictionary
        :return: the dictionaries, one for each translation
        """
        raise NotImplementedError

    def _create_token_types_graph(self) -> Dict[str, List[str]]:
        """ Creates a dictionary for the directions of the token types of the encoding"""
        raise NotImplementedError

    def _create_durations_tuples(self) -> List[Tuple]:
        """ Creates the possible durations in bar / beat units, as tuple of the form:
        (beat, pos, res) where beat is the number of beats, pos the number of "frames"
        ans res the beat resolution considered (frames per beat)
        Example: (2, 5, 8) means the duration is 2 beat long + position 5 / 8 of the ongoing beat
        In pure ticks we have: duration = (beat * res + pos) * time_division // res
            Is equivalent to: duration = nb_of_frames * ticks_per_frame
        So in the last example, if time_division is 384: duration = (2 * 8 + 5) * 384 // 8 = 1008 ticks

        :return: the duration bins
        """
        durations = []
        for beat_range, beat_res in self.beat_res.items():
            durations += [(beat, pos, beat_res) for beat in range(*beat_range) for pos in range(beat_res)]
        durations += [(max(max(self.beat_res)), 0, self.beat_res[max(self.beat_res)])]  # the last one
        del durations[0]  # removes duration of 0
        return durations

    def tokenize_midi_dataset(self, midi_paths: Union[List[str], List[Path], List[PurePath]],
                              out_dir: Union[str, Path, PurePath], validation_fn: Callable[[MidiFile], bool] = None,
                              logging: bool = True):
        """ Converts a dataset / list of MIDI files, into their token version and save them as json files

        :param midi_paths: paths of the MIDI files
        :param out_dir: output directory to save the converted files
        :param validation_fn: a function checking if the MIDI is valid on your requirements
                            (e.g. time signature, minimum/maximum length, instruments ...)
        :param logging: logs a progress bar
        """
        Path(out_dir).mkdir(parents=True, exist_ok=True)

        # Making a directory of the parent folders for the JSON file
        # parent_dir = PurePath(midi_paths[0]).parent[0]
        # PurePath(out_dir, parent_dir).mkdir(parents=True, exist_ok=True)

        for m, midi_path in enumerate(midi_paths):
            if logging:
                bar_len = 60
                filled_len = int(round(bar_len * m / len(midi_paths)))
                percents = round(100.0 * m / len(midi_paths), 2)
                bar = '=' * filled_len + '-' * (bar_len - filled_len)
                prog = f'\r{m} / {len(midi_paths)} [{bar}] {percents:.1f}% ...Converting MIDIs to tokens: {midi_path}'
                stdout.write(prog)
                stdout.flush()

            # Some MIDIs can contains errors that are raised by Mido, if so the loop continues
            try:
                midi = MidiFile(PurePath(midi_path))
            except Exception as _:  # ValueError, OSError, FileNotFoundError, IOError, EOFError, mido.KeySignatureError
                continue

            # Passing the MIDI to validation tests if given
            if validation_fn is not None:
                if not validation_fn(midi):
                    continue

            # Converting the MIDI to tokens and saving them as json
            tokens = self.midi_to_tokens(midi)
            midi_name = PurePath(midi_path).stem
            with open(PurePath(out_dir, midi_name).with_suffix(".json"), 'w') as outfile:
                json.dump([tokens[0], track_info[0]], outfile)

        self.save_params(out_dir)  # Saves the parameters with which the MIDIs are converted

    def save_params(self, out_dir: Union[str, Path, PurePath]):
        """ Saves the base parameters of this encoding in a txt file
        Useful to keep track of how a dataset has been tokenized / encoded
        It will also save the name of the class used, i.e. the encoding strategy

        :param out_dir: output directory to save the file
        """
        Path(out_dir).mkdir(parents=True, exist_ok=True)
        with open(PurePath(out_dir, 'config').with_suffix(".txt"), 'w') as outfile:
            json.dump({'pitch_range': (self.pitch_range.start, self.pitch_range.stop),
                       'beat_res': self.beat_res, 'nb_velocities': len(self.velocity_bins),
                       'additional_tokens': self.additional_tokens,
                       'encoding': self.__class__.__name__}, outfile)

    def load_params(self, params: Union[str, Path, PurePath, Dict[str, Any]]):
        """ Load parameters and set the encoder attributes

        :param params: can be a path to the parameter (json encoded) file or a dictionary
        """
        if isinstance(params, (str, Path, PurePath)):
            with open(params) as param_file:
                params = json.load(param_file)

        if not isinstance(params['pitch_range'], range):
            params['pitch_range'] = range(*params['pitch_range'])

        for key, value in params.items():
            setattr(self, key, value)

def get_midi_programs(midi: MidiFile) -> List[Tuple[int, bool]]:
    """ Returns the list of programs of the tracks of a MIDI, deeping the
    same order. It returns it as a list of tuples (program, is_drum).

    :param midi: the MIDI object to extract tracks programs
    :return: the list of track programs, as a list of tuples (program, is_drum)
    """
    return [(int(track.program), track.is_drum) for track in midi.instruments]

def quantize_note_times(notes: List[Note], time_division: int, beat_res: int):
    """ Quantize the notes items start and end values.
    It shifts the notes so they start at times that match the quantization (e.g. 16 frames per bar)

    :param notes: notes to quantize
    :param time_division: MIDI time division / resolution, in ticks/beat (of the MIDI being parsed)
    :param beat_res: number of frames (time steps, or positions) per beat
    """
    ticks = int(time_division / beat_res)
    quantized_ticks = np.arange(0, max([n.end for n in notes]) + 2 * ticks, ticks, dtype=int)
    for i, note in enumerate(notes):  # items are notes
        note.start = quantized_ticks[np.argmin(np.abs(quantized_ticks - note.start))]
        note.end = quantized_ticks[np.argmin(np.abs(quantized_ticks - note.end))]

        if note.start == note.end:  # if this happens to often, consider using a higher beat resolution
            note.end += ticks  # like 8 frames per beat or 24 frames per bar

def quantize_tempos(tempos: List[TempoChange], time_division: int, beat_res: int):
    """ Quantize the times of tempo change events.

    :param tempos: tempo changes to quantize
    :param time_division: MIDI time division / resolution, in ticks/beat (of the MIDI being parsed)
    :param beat_res: number of frames (time steps, or positions) per beat
    """
    ticks = int(time_division / beat_res)
    quantized_ticks = np.arange(0, max([t.time for t in tempos]) + 2 * ticks, ticks, dtype=int)
    for tempo_change in tempos:
        tempo_change.time = quantized_ticks[np.argmin(np.abs(quantized_ticks - tempo_change.time))]

def remove_duplicated_notes(notes: List[Note]):
    """ Remove possible duplicated notes, i.e. with the same pitch, starting and ending times.
    Before running this function make sure the notes has been sorted by start and pitch:
    notes.sort(key=lambda x: (x.start, x.pitch))

    :param notes: notes to analyse
    """
    for i in range(len(notes) - 1, 0, -1):  # removing possible duplicated notes
        if notes[i].pitch == notes[i - 1].pitch and notes[i].start == notes[i - 1].start and \
                notes[i].end == notes[i - 1].end:
            del notes[i]

def detect_chords(notes: List[Note], time_division: int) -> List[Event]:
    """ Chord detection method.
    NOTE: on very large tracks with high note density this method can be very slow !
    If you plan to use it with the Maestro or GiantMIDI datasets, it can take up to
    hundreds of seconds per MIDI depending on your cpu.
    One time step at a time, it will analyse the notes played together
    and detect possible chords.

    :param notes: notes to analyse
    :param time_division: MIDI time division / resolution, in ticks/beat (of the MIDI being parsed)
    :return: the detected chords as Event objects
    """
    tuples = []
    for note in notes:
        tuples.append((note.pitch, int(note.start), int(note.end)))
    notes = np.asarray(tuples)

    count = 0
    chords = []
    while count < len(notes):
        # Gather the notes around the same time step
        onset_notes = notes[count:]
        onset_notes = onset_notes[np.where(onset_notes[:, 1] <= notes[count][1] + time_division / 4)]

        # If it is ambiguous, e.g. the notes lengths are too different
        if np.any(np.abs(onset_notes[:, 2] - onset_notes[0, 2]) > time_division / 2):
            count += len(onset_notes)
            continue

        # Selects the possible chords notes
        if notes[count][2] - notes[count][1] <= time_division / 2:
            onset_notes = onset_notes[np.where(onset_notes[:, 1] == onset_notes[0][1])]
        chord = onset_notes[np.where(onset_notes[:, 2] - onset_notes[0, 2] <= time_division / 2)]

        # Creates the "chord map" and see if it has a "known" quality, append a chord event if it is valid
        chord_map = (chord[:, 0] - chord[0, 0]).tolist()
        if 3 <= len(chord_map) <= 5 and chord_map[-1] <= 24:  # max interval between the root and highest degree
            chord_quality = len(chord)
            for quality, known_chord in CHORD_MAPS.items():
                if known_chord == chord_map:
                    chord_quality = quality
                    break
            chords.append((chord_quality, min(chord[:, 1]), chord_map))
        count += len(onset_notes)  # Move to the next notes

    events = []
    for chord in chords:
        events.append(Event('Chord', chord[1], chord[0], chord[2]))
    return events

def _detect_chords_python(notes: List[Note], time_division: int) -> List[Event]:
    """ DEPRECIATED
    Old chord detection method, equivalent to detect_chords but 100% python
    The code is more elegant but slower

    :param notes: notes to analyse
    :param time_division: MIDI time division / resolution, in ticks/beat (of the MIDI being parsed)
    :return: the detected chords as Event objects
    """
    count = 0
    chords = []
    while count < len(notes):
        # Gather the notes around the same time step
        # onset_notes = [note for note in notes if abs(note.start - notes[count].start) <= time_division / 4]
        onset_notes = [note for note in notes[count:] if note.start - notes[count].start <= time_division / 4]

        # If it is ambiguous, e.g. the notes lengths are too different
        if any(abs(note.end - notes[count].end) > time_division / 2 for note in onset_notes):
            count += len(onset_notes)
            continue

        # Selects the possible chords notes
        if notes[count].end - notes[count].start <= time_division / 2:
            onset_notes = [note for note in notes[count:] if note.start == notes[count].start]
        chord = [note for note in onset_notes if abs(note.end - notes[count].end) <= time_division / 2]

        # Creates the "chord map" and see if it has a "known" quality, append a chord event if it is valid
        chord_map = [chord[i].pitch - chord[0].pitch for i in range(len(chord))]
        if 3 <= len(chord_map) <= 5 and chord_map[-1] <= 24:  # max interval between the root and highest degree
            if chord_map in CHORD_MAPS.values():
                chord_quality = list(CHORD_MAPS.keys())[list(CHORD_MAPS.values()).index(chord_map)]
                chords.append(Event('Chord', min(note.start for note in chord), chord_quality, chord_map))
            else:
                chords.append(Event('Chord', min(note.start for note in chord), len(chord), chord_map))
        count += len(onset_notes)  # Move to the next notes
    return chords

def merge_tracks(tracks: List[Instrument]) -> Instrument:
    """ Merge several miditoolkit Instrument objects
    It will take the first object, and concat the notes of the others

    :param tracks: list of tracks to merge
    :return: the merged track
    """
    tracks[0].name += ''.join([' / ' + t.name for t in tracks[1:]])
    tracks[0].notes = sum((t.notes for t in tracks), [])
    tracks[0].notes.sort(key=lambda note: note.start)
    return tracks[0]

def current_bar_pos(seq: List[int], bar_token: int, position_tokens: List[int], pitch_tokens: List[int],
                    chord_tokens: List[int] = None) -> Tuple[int, int, List[int], bool]:
    """ Detects the current state of a sequence of tokens

    :param seq: sequence of tokens
    :param bar_token: the bar token value
    :param position_tokens: position tokens values
    :param pitch_tokens: pitch tokens values
    :param chord_tokens: chord tokens values
    :return: the current bar, current position within the bar, current pitches played at this position,
            and if a chord token has been predicted at this position
    """
    # Current bar
    bar_idx = [i for i, token in enumerate(seq) if token == bar_token]
    current_bar = len(bar_idx)
    # Current position value within the bar
    pos_idx = [i for i, token in enumerate(seq[bar_idx[-1]:]) if token in position_tokens]
    current_pos = len(pos_idx) - 1  # position value, e.g. from 0 to 15, -1 means a bar with no Pos token following
    # Pitches played at the current position
    current_pitches = [token for token in seq[pos_idx[-1]:] if token in pitch_tokens]
    # Chord predicted
    if chord_tokens is not None:
        chord_at_this_pos = any(token in chord_tokens for token in seq[pos_idx[-1]:])
    else:
        chord_at_this_pos = False
    return current_bar, current_pos, current_pitches, chord_at_this_pos

ext: py

reason: import numpy


text:
""" Script used for the main functionnalities of the woods package

There is 2 mode of operation:
    - training mode: trains a model on a given dataset with a given test environment using a given algorithm
    - test mode: tests an existing model on a given dataset with a given test environment using a given algorithm

Raises:
    NotImplementedError: Some part of the code is not implemented yet
"""

import os
import json
import time
import random
import argparse
import numpy as np

import torch
from torch import nn, optim

from woods import datasets
from woods import models
from woods import objectives
from woods import hyperparams
from woods import utils
from woods.train import train, get_accuracies

if __name__ == '__main__':

    ## Args
    parser = argparse.ArgumentParser(description='Train a model on a dataset with an objective and test on a test_env')
    # Main mode
    parser.add_argument('mode', choices=['train', 'eval'])
    # Dataset arguments
    parser.add_argument('--test_env', type=int, default = None)
    parser.add_argument('--dataset', type=str)
    parser.add_argument('--holdout_fraction', type=float, default=0.2)
    # Setup arguments
    parser.add_argument('--objective', type=str, choices=objectives.OBJECTIVES)
    # Hyperparameters arguments
    parser.add_argument('--sample_hparams', action='store_true')
    parser.add_argument('--hparams_seed', type=int, default=0, help='Seed for random hparams (Is not used if sample_hparams is not true')
    parser.add_argument('--trial_seed', type=int, default=0, help='Trial number for seeding split_dataset and random_hparams.')
    parser.add_argument('--seed', type=int, default=0, help='Seed for everything else')
    # Directory arguments
    parser.add_argument('--data_path', type=str, default='~/Documents/Data/')
    parser.add_argument('--save_path', type=str, default='./results/')
    parser.add_argument('--download', action='store_true')
    # Model evaluation arguments
    parser.add_argument('--save', action='store_true')
    parser.add_argument('--model_path', type=str, default=None)

    flags = parser.parse_args()

    # Device definition
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print("Using CUDA device ", torch.cuda.current_device())
    else:
        device = torch.device("cpu")

    print('Flags:')
    for k,v in sorted(vars(flags).items()):
        print("\t{}: {}".format(k, v))

    ## Making job ID and checking if done
    job_name = utils.get_job_name(vars(flags))

    assert isinstance(flags.test_env, int) or flags.test_env is None, "Invalid test environment"
    if flags.mode == 'train':
        assert not os.path.isfile(os.path.join(flags.save_path, 'logs', job_name+'.json')), "\n*********************************\n*** Job Already ran and saved ***\n*********************************\n"

    ## Getting hparams
    hparam_sampling_seed = utils.seed_hash(flags.hparams_seed, flags.trial_seed)
    training_hparams = hyperparams.get_training_hparams(flags.dataset, hparam_sampling_seed, flags.sample_hparams)
    training_hparams['device'] = device
    objective_hparams = hyperparams.get_objective_hparams(flags.objective, hparam_sampling_seed, flags.sample_hparams)
    objective_hparams['device'] = device
    model_hparams = hyperparams.get_model_hparams(flags.dataset)
    model_hparams['device'] = device
    model_hparams['model_path'] = flags.model_path

    print('HParams:')
    for k, v in sorted(training_hparams.items()):
        print('\t{}: {}'.format(k, v))
    for k, v in sorted(model_hparams.items()):
        print('\t{}: {}'.format(k, v))
    for k, v in sorted(objective_hparams.items()):
        print('\t{}: {}'.format(k, v))

    ## Make dataset
    dataset_class = datasets.get_dataset_class(flags.dataset)
    dataset = dataset_class(flags, training_hparams)
    _, in_loaders = dataset.get_train_loaders()

    # Make some checks about the dataset
    if datasets.num_environments(flags.dataset) == 1:
        assert flags.objective == 'ERM', "Dataset has only one environment, cannot compute multi-environment penalties"

    ## Setting global seed
    random.seed(flags.seed)
    np.random.seed(flags.seed)
    torch.manual_seed(flags.seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    ## Initialize a model to train
    model = models.get_model(dataset, model_hparams)
    print("Number of parameters = ", sum(p.numel() for p in model.parameters() if p.requires_grad))

    # Define training aid
    parameters = [p for p in model.parameters() if p.requires_grad]
    optimizer = optim.Adam(parameters, lr=training_hparams['lr'], weight_decay=training_hparams['weight_decay'])

    ## Initialize some Objective
    objective_class = objectives.get_objective_class(flags.objective)
    objective = objective_class(model, dataset, optimizer, objective_hparams)

    ## Do the thing
    model.to(device)
    if flags.mode == 'train':

        model, record, table = train(flags, training_hparams, model, objective, dataset, device)

        ## Save stuff
        if flags.save:
            hparams = {}
            del training_hparams['device']
            del model_hparams['device']
            del objective_hparams['device']
            hparams.update(training_hparams)
            hparams.update(model_hparams)
            hparams.update(objective_hparams)
            record['hparams'] = hparams
            record['flags'] = vars(flags)
            os.makedirs(os.path.join(flags.save_path, 'logs'), exist_ok=True)
            with open(os.path.join(flags.save_path, 'logs', job_name+'.json'), 'w') as f:
                json.dump(record, f)
            os.makedirs(os.path.join(flags.save_path, 'models'), exist_ok=True)
            torch.save(model.state_dict(), os.path.join(flags.save_path, 'models', job_name+'.pt'))
            os.makedirs(os.path.join(flags.save_path, 'outputs'), exist_ok=True)
            with open(os.path.join(flags.save_path, 'outputs', job_name+'.txt'), 'w') as f:
                f.write('HParams:\n')
                for k, v in sorted(training_hparams.items()):
                    f.write('\t{}: {}\n'.format(k, v))
                for k, v in sorted(model_hparams.items()):
                    f.write('\t{}: {}\n'.format(k, v))
                for k, v in sorted(objective_hparams.items()):
                    f.write('\t{}: {}\n'.format(k, v))
                job_id = 'Training ' + flags.objective  + ' on ' + flags.dataset + ' (H=' + str(flags.hparams_seed) + ', T=' + str(flags.trial_seed) + ')'
                f.write(table.get_string(title=job_id, border=True, hrule=0))

    elif flags.mode == 'eval':
        # raise NotImplementedError('This part is in quarantine')
        # Load the weights
        assert flags.model_path != None, "You must give the model_path in order to evaluate a model"
        model.load_state_dict(torch.load(os.path.join(flags.model_path)))

        # Get accuracies
        val_start = time.time()
        record = get_accuracies(objective, dataset, device)
        val_time = time.time() - val_start

        # train_names, _ = dataset.get_train_loaders()
        # t = utils.setup_pretty_table(flags)
        # if dataset.TASK == 'regression':
        #     t.add_row(['eval']
        #             + ["{:.1e} :: {:.1e}".format(record[str(e)+'_in_loss'], record[str(e)+'_out_loss']) for e in dataset.ENVS]
        #             + ["{:.1e}".format(np.average([record[str(e)+'_loss'] for e in train_names]))]
        #             + ['.']
        #             + ['.']
        #             + ["{:.2f}".format(val_time)])
        # else:
        #     t.add_row(['eval']
        #             + ["{:.2f} :: {:.2f}".format(record[str(e)+'_in_acc'], record[str(e)+'_out_acc']) for e in dataset.ENVS]
        #             + ["{:.2f}".format(np.average([record[str(e)+'_loss'] for e in train_names]))]
        #             + ['.']
        #             + ['.']
        #             + ["{:.2f}".format(val_time)])
        # print("\n".join(t.get_string().splitlines()[-2:]))

max_stars_repo_name: jc-audet/WOODS

max_issues_repo_name: jc-audet/WOODS

max_forks_repo_name: jc-audet/WOODS


text:
import Tkinter as Tk
import scipy.signal as sig
import numpy as np
from scipy.io import wavfile
import matplotlib.pyplot as pl

class EQModel(object):
    def __init__(self, signal_in, gains):
        self.signal = signal_in
        self.gains = gains
        self.Q = 3.5

    def apply_gains(self):
        """From signals and systems for dummies"""
        fs = 44100.0
        NB = len(self.gains)
        fc = 31.25*2**np.arange(5)
        A = np.zeros((NB, 3))
        B = np.zeros((NB, 3))

        for k in range(NB):
            [b, a] = self.peaking(self.gains[k], fc[k], self.Q)
            B[k,:] = b
            A[k,:] = a

        y = np.zeros(len(self.signal))
        for k in range(NB):
            if k == 0:
                y = sig.lfilter(B[k,:], A[k,:], self.signal)
            else:
                y = sig.lfilter(B[k,:], A[k,:], y)
        return y

    def peaking(self, gain, fc, Q=3.5, fs=44100.):
        """From Signals and Systems for dummies"""
        mu = 10**(gain/20.)
	kq = 4/(1 + mu)*np.tan(2*np.pi*fc/fs/(2*Q))
	Cpk = (1 + kq *mu)/(1 + kq)
	b1 = -2*np.cos(2*np.pi*fc/fs)/(1 + kq*mu)
	b2 = (1 - kq*mu)/(1 + kq*mu)
	a1 = -2*np.cos(2*np.pi*fc/fs)/(1 + kq)
	a2 = (1 - kq)/(1 + kq)
	b = Cpk*np.array([1, b1, b2])
	a = np.array([1, a1, a2])
	return b,a

class EQView(object):
    def __init__(self, master):
        self.frame = Tk.Frame(master)
        self.guiElements = {}
        for i in range(0, 5):
            self.guiElements['Scale' + str(i)] = Tk.Scale(master)
        self.guiElements['Scale0'].grid(row=0, column=0)
        self.guiElements['Scale1'].grid(row=0, column=1)
        self.guiElements['Scale2'].grid(row=0, column=2)
        self.guiElements['Scale3'].grid(row=0, column=3)
        self.guiElements['Scale4'].grid(row=0, column=4)
        self.button = Tk.Button(master, text="Apply filters")
        self.button.grid(row=1, column=0)

class EQController(object):
    def __init__(self):
        self.root = Tk.Tk()
        self.t = np.arange(0, 0.05, 1/44100.0)
        self.x = np.cos(2*np.pi*31.25*self.t)+np.cos(2*np.pi*4000*self.t)
        self.gains = np.zeros(5)
        self.model = EQModel(self.x, self.gains)
        self.view = EQView(self.root)
        self.view.button.config(command=self.apply_gains)

    def run(self):
        self.root.deiconify()
        self.root.mainloop()

    def apply_gains(self):
        for i in range(0, 5):
            self.model.gains[i] = self.view.guiElements['Scale' + str(i)].get()
        out =self.model.apply_gains()
        pl.plot(self.t, self.x, label='orig')
        pl.plot(self.t, out, label='equalised')
        pl.legend()
        pl.show()

def main():
    app = EQController()
    app.run()

if __name__ == '__main__':
    main()

reason: import numpy,import scipy,from scipy


text:
#!/usr/bin/python3

## ================================================================================
## Copyright: (C) 2022, SIRRL Social and Intelligent Robotics Research Laboratory,
##     University of Waterloo, All rights reserved.
##
## Authors:
##     Austin Kothig <austin.kothig@uwaterloo.ca>
##
## CopyPolicy: Released under the terms of the MIT License.
##     See the accompanying LICENSE file for details.
## ================================================================================

import numpy as np
import cv2
import time

import yarp

import argparse

def getArgs():
    parser = argparse.ArgumentParser(description='yarpMediaPlayer')
    parser.add_argument('-n', '--name',    default='/mediaPlayer',     help='Name for the module.          (default: {})'.format('/mediaPlayer'))
    parser.add_argument('-x',              default=0,     type=int,    help='Position for window to be.    (default: {})'.format(0))
    parser.add_argument('-y',              default=0,     type=int,    help='Position for window to be.    (default: {})'.format(0))
    parser.add_argument('-H',              default=540,   type=int,    help='Height for the window to be.  (default: {})'.format(540))
    parser.add_argument('-W',              default=960,   type=int,    help='Width for the window to be.   (default: {})'.format(960))
    parser.add_argument('-d', '--default', default=None,               help='Default video to play.        (default: {})'.format(None))
    parser.add_argument('-g', '--goal',    default=None,               help='Goal image to send.           (default: {})'.format(None))
    parser.add_argument('-s', '--speed',   default=10.,   type=float,  help='Speed mutl when jobs todo.    (default: {})'.format(10.))
    parser.add_argument('-b', '--breako',  default=False,              help='Allow ESC and Q to break out? (default: {})'.format(False))
    args = parser.parse_args()
    return args

# Init yarp network.
yarp.Network.init()

class mediaPlayer(object):
    '''
    Receive requests to open videos. Play them through opencv.
    '''
    def __init__(self, args):

        # Get some args.
        self.name       = args.name
        self.default    = args.default
        self.goal       = args.goal
        self.speedup    = args.speed
        self.break_outs = args.breako
        self.xpos       = args.x
        self.ypos       = args.y
        self.height     = args.H
        self.width      = args.W

        self.file_buffer = []

        # Open the audio and video streams.
        self.video = cv2.VideoCapture(self.default)
        self.fps = self.video.get(cv2.CAP_PROP_FPS)

        if self.goal != None:
            goal_img = cv2.imread(self.goal)
            cv2.namedWindow('goal state')
            cv2.imshow('goal state', goal_img)
            #cv2.moveWindow('goal state', self.height//3, self.width)
            cv2.moveWindow('goal state', self.width, self.height//3)

        # Init a flag for completion of current media.
        self.playing = True

        # Open the yarp port.
        self.port = yarp.BufferedPortBottle()
        self.port.open(self.name + ":i")

        return

    def run(self):

        while True:

            self.checkPort()

            if self.playing:

                while self.video.isOpened():

                    ret, frame = self.video.read()

                    if not ret:
                        # End of video.
                        self.playing = False
                        self.video.release()
                        break

                    # Wait based on the frames-per-second of the video,
                    # but if we have multiple videos in the queue, speed
                    # things up so that we're not waiting too for the
                    # most recent video clips.
                    wait_for = int((1000 / self.fps) / (self.speedup if len(self.file_buffer) > 1 else 1.))
                    key = cv2.waitKey(wait_for)

                    # Allow breaking out?
                    if (key == 27 or key == 1048603) and self.break_outs:
                        print("Broke out?")
                        self.cleanup()
                        exit(0)
                        break

                    # Modify the frame.
                    frame = cv2.resize(frame,
                        (self.width, self.height),
                        fx=0, fy=0,
                        interpolation=cv2.INTER_CUBIC
                    )

                    cv2.imshow(self.name, frame)
                    cv2.moveWindow(self.name, self.xpos, self.ypos)

                    self.checkPort()

            # If there is something in the beffer, load it.
            elif len(self.file_buffer):

                # Get the file at the front.
                file = self.file_buffer[0]
                self.file_buffer.pop(0)

                if file == "none":
                    file = self.default
                    if file == None: continue

                # Load the file.
                print("Loading file: {}".format(file))
                self.video.open(file)

                if not self.video.isOpened():
                    print("Could not open ``{}``!!".format(file))

                else:
                    self.fps = self.video.get(cv2.CAP_PROP_FPS)
                    self.playing = True

            # Else wait for some input.
            else:
                print("Waiting for video . . .")
                time.sleep(0.1)

        return

    def checkPort(self):

        # See if we have mail.
        b = self.port.read(False)
        if b is None:
            return

        # Get the contents of the bottle.
        new_file = b.get(0).toString()
        if new_file == "none" and new_file in self.file_buffer:
            # Ensure only one copy of "none" in buffer.
            return

        # Check if we're told to quit.
        if new_file == "exit":
            self.cleanup()
            exit(0)

        # Add this video to the queue.
        self.file_buffer.append(new_file)

        return

    def cleanup(self):

        print("Closing media streams.")
        if self.video.isOpened():
            self.video.release()
            cv2.destroyAllWindows()

        print("Closing yarp ports.")
        self.port.interrupt()
        self.port.close()
        return

def main():

    # Parse the arguments.
    args = getArgs()

    player = mediaPlayer(args)

    try:
        player.run()
    finally:
        player.cleanup()

    return

if __name__ == '__main__':
    main()

max_stars_repo_name: kothiga/Embodied-Social-Interface

max_issues_repo_name: kothiga/Embodied-Social-Interface

max_forks_repo_name: kothiga/Embodied-Social-Interface


text:
import numpy as np, math
from threading import RLock
import sys, copy
import matplotlib.pyplot as pp

import roslib; roslib.load_manifest('hrl_software_simulation_darpa_m3')

import rospy
import PyKDL as kdl

from equilibrium_point_control.hrl_arm import HRLArm, HRLArmKinematics

import hrl_lib.geometry as hg
import hrl_lib.kdl_utils as ku
import hrl_lib.viz as hv

from hrl_msgs.msg import FloatArrayBare
from visualization_msgs.msg import Marker
from hrl_haptic_manipulation_in_clutter_msgs.msg import MechanicalImpedanceParams

class ODESimArm(HRLArm):
    def __init__(self, d_robot):
        rospy.loginfo("Loading ODESimArm")
        kinematics = RobotSimulatorKDL(d_robot) # KDL chain.
        HRLArm.__init__(self, kinematics)

        self.jep_pub = rospy.Publisher('/sim_arm/command/jep', FloatArrayBare)
        self.cep_marker_pub = rospy.Publisher('/sim_arm/viz/cep', Marker)

        self.impedance_pub = rospy.Publisher('/sim_arm/command/joint_impedance',
                                             MechanicalImpedanceParams)

        rospy.Subscriber('/sim_arm/joint_angles', FloatArrayBare,
                         self.joint_states_cb)
        rospy.Subscriber('/sim_arm/joint_angle_rates', FloatArrayBare,
                         self.joint_rates_cb)
        rospy.Subscriber('/sim_arm/jep', FloatArrayBare, self.jep_cb)
        rospy.Subscriber('/sim_arm/joint_impedance', MechanicalImpedanceParams,
                         self.impedance_params_cb)
        rospy.Subscriber ("/delta_jep_mpc_cvxgen", FloatArrayBare, self.set_delta_ep_ros)

        rospy.sleep(1.)
        rospy.loginfo("Finished loading SimpleArmManger")

    def jep_cb(self, msg):
        with self.lock:
            self.ep = copy.copy(msg.data)

    def joint_states_cb(self, data):
        with self.lock:
            self.q = copy.copy(data.data)

    def joint_rates_cb(self, data):
        with self.lock:
            self.qdot = copy.copy(data.data)

    def impedance_params_cb(self, data):
        with self.lock:
            self.kp = copy.copy(data.k_p.data)
            self.kd = copy.copy(data.k_d.data)

    def get_joint_velocities(self):
        with self.lock:
            qdot = copy.copy(self.qdot)
        return qdot

    def set_ep(self, jep, duration=0.15):
        f = FloatArrayBare(jep)
        self.jep_pub.publish(f)
        self.publish_rviz_markers()

    def set_delta_ep_ros(self, msg):
        delta_jep = copy.copy(msg.data)
        if delta_jep is None or len(delta_jep) != 3:
            raise RuntimeError("set_jep value is " + str(delta_jep))

        with self.lock:
          if self.ep == None:
            self.ep = self.get_joint_angles()

          jep = (np.array(self.ep) + np.array(delta_jep)).tolist()

          f = FloatArrayBare(jep)
          self.jep_pub.publish(f)

    def publish_rviz_markers(self):
        # publish the CEP marker.
        jep = self.get_ep()
        cep, r = self.kinematics.FK(jep)
        o = np.matrix([0.,0.,0.,1.]).T
        cep_marker = hv.single_marker(cep, o, 'sphere',
                        '/torso_lift_link', color=(0., 0., 1., 1.),
                        scale = (0.02, 0.02, 0.02), duration=0.)

        cep_marker.header.stamp = rospy.Time.now()
        self.cep_marker_pub.publish(cep_marker)

    # kp, kd - list or np array of stiffness and damping.
    def set_joint_impedance(self, kp, kd):
        im = MechanicalImpedanceParams()
        im.k_p.data = kp
        im.k_d.data = kd
        self.impedance_pub.publish(im)

##
# KDL for kinematics etc.
class RobotSimulatorKDL(HRLArmKinematics):
    def __init__(self, d_robot):
        HRLArmKinematics.__init__(self, len(d_robot.b_jt_anchor))
        self.arm_type = "simulated"
        self.d_robot = d_robot
        self.right_chain = self.create_right_chain()
        fk, ik_v, ik_p, jac = self.create_solvers(self.right_chain)

        self.right_fk = fk
        self.right_ik_v = ik_v
        self.right_ik_p = ik_p
        self.right_jac = jac
        self.right_tooltip = np.matrix([0.,0.,0.]).T

        # marc joint limits.
        self.min_jtlim_arr = d_robot.b_jt_limits_min
        self.max_jtlim_arr = d_robot.b_jt_limits_max

    def create_right_chain(self):
        height = 0.0
        linkage_offset_from_ground = np.matrix([0., 0., height]).T
        self.linkage_offset_from_ground = linkage_offset_from_ground
        self.n_jts = len(self.d_robot.b_jt_anchor)
        rospy.loginfo("number of joints is :"+str(self.n_jts))

        ee_location = self.d_robot.ee_location
        b_jt_anchor = self.d_robot.b_jt_anchor
        b_jt_axis = self.d_robot.b_jt_axis

        ch = kdl.Chain()
        prev_vec = np.copy(linkage_offset_from_ground.A1)
        n = len(self.d_robot.b_jt_anchor)

        for i in xrange(n-1):
            if b_jt_axis[i][0] == 1 and b_jt_axis[i][1] == 0 and b_jt_axis[i][2] == 0:
                kdl_jt = kdl.Joint(kdl.Joint.RotX)
            elif b_jt_axis[i][0] == 0 and b_jt_axis[i][1] == 1 and b_jt_axis[i][2] == 0:
                kdl_jt = kdl.Joint(kdl.Joint.RotY)
            elif b_jt_axis[i][0] == 0 and b_jt_axis[i][1] == 0 and b_jt_axis[i][2] == 1:
                kdl_jt = kdl.Joint(kdl.Joint.RotZ)
            else:
                print "can't do off-axis joints yet!!!"

            np_vec = np.array(b_jt_anchor[i+1])
            diff_vec = np_vec-prev_vec
            prev_vec = np_vec
            kdl_vec = kdl.Vector(diff_vec[0], diff_vec[1], diff_vec[2])
            ch.addSegment(kdl.Segment(kdl_jt, kdl.Frame(kdl_vec)))

        np_vec = np.copy(ee_location.A1)
        diff_vec = np_vec-prev_vec

        if b_jt_axis[n-1][0] == 1 and b_jt_axis[n-1][1] == 0 and b_jt_axis[n-1][2] == 0:
            kdl_jt = kdl.Joint(kdl.Joint.RotX)
        elif b_jt_axis[n-1][0] == 0 and b_jt_axis[n-1][1] == 1 and b_jt_axis[n-1][2] == 0:
            kdl_jt = kdl.Joint(kdl.Joint.RotY)
        elif b_jt_axis[n-1][0] == 0 and b_jt_axis[n-1][1] == 0 and b_jt_axis[n-1][2] == 1:
            kdl_jt = kdl.Joint(kdl.Joint.RotZ)
        else:
            print "can't do off-axis joints yet!!!"

        kdl_vec = kdl.Vector(diff_vec[0], diff_vec[1], diff_vec[2])
        ch.addSegment(kdl.Segment(kdl_jt, kdl.Frame(kdl_vec)))

        return ch

    def create_solvers(self, ch):
         fk = kdl.ChainFkSolverPos_recursive(ch)
         ik_v = kdl.ChainIkSolverVel_pinv(ch)
         ik_p = kdl.ChainIkSolverPos_NR(ch, fk, ik_v)
         jac = kdl.ChainJntToJacSolver(ch)
         return fk, ik_v, ik_p, jac

    def FK_kdl(self, q, link_number):
        fk = self.right_fk
        endeffec_frame = kdl.Frame()
        kinematics_status = fk.JntToCart(q, endeffec_frame,
                                         link_number)
        if kinematics_status >= 0:
            return endeffec_frame
        else:
            rospy.loginfo('Could not compute forward kinematics.')
            return None

    ## returns point in torso lift link.
    def FK_vanilla(self, q, link_number=None):
        if link_number == None:
            link_number = self.n_jts
        link_number = min(link_number, self.n_jts)

        q = self.list_to_kdl(q)
        frame = self.FK_kdl(q, link_number)
        pos = frame.p
        pos = ku.kdl_vec_to_np(pos)
        pos = pos + self.linkage_offset_from_ground
        m = frame.M
        rot = ku.kdl_rot_to_np(m)

        return pos, rot

    def kdl_to_list(self, q):
        if q == None:
            return None
        n = self.n_jts
        q_list = [0] * n
        for i in xrange(n):
            q_list[i] = q[i]
        return q_list

    def list_to_kdl(self, q):
        if q == None:
            return None
        n = len(q)
        q_kdl = kdl.JntArray(n)
        for i in xrange(n):
            q_kdl[i] = q[i]
        return q_kdl

    ## compute Jacobian at point pos.
    # p is in the ground coord frame.
    def jacobian(self, q, pos=None):
        if pos == None:
            pos = self.FK(q)[0]

        v_list = []
        w_list = []

        #ee, r = self.FK(q)
        #pos = ee

        for i in xrange(self.n_jts):
            p, rot = self.FK_vanilla(q, i)
            r = pos - p
            z_idx = self.right_chain.getSegment(i).getJoint().getType() - 1
            z = rot[:, z_idx]
            v_list.append(np.matrix(np.cross(z.A1, r.A1)).T)
            w_list.append(z)

        J = np.row_stack((np.column_stack(v_list), np.column_stack(w_list)))
        return J

    def clamp_to_joint_limits(self, q):
        return np.clip(q, self.min_jtlim_arr, self.max_jtlim_arr)

    def within_joint_limits(self, q, delta_list=None):
        if delta_list == None:
            delta_list = [0]*len(q)

        min_arr = self.min_jtlim_arr
        max_arr = self.max_jtlim_arr

        q_arr = np.array(q)
        d_arr = np.array(delta_list)

        return np.all((q_arr <= max_arr+d_arr, q_arr >= min_arr+d_arr))

    def get_joint_limits(self):
        return self.min_jtlim_arr, self.max_jtlim_arr

    # plot the arm using matplotlib.
    # flip_xy - x will point up and +ve y will be to the left. This is
    # to match how we look at the arm in rviz.
    def plot_arm(self, q, color, alpha, flip_xy, linewidth=2):
        pts = [[0.,0.,0.]]
        for i in range(len(q)):
            p,_ = self.FK(q, i+1)
            pts.append(p.A1.tolist())

        pts_2d = np.array(pts)[:,0:2]
        direc_list = (pts_2d[1:] - pts_2d[:-1]).tolist()

        for i, d in enumerate(direc_list):
            d_vec = np.matrix(d).T
            d_vec = d_vec / np.linalg.norm(d_vec)
            w = np.cross(d_vec.A1, np.array([0., 0., 1.])) * 0.03/2
            x1 = pts_2d[i,0]
            y1 = pts_2d[i,1]
            x2 = pts_2d[i+1,0]
            y2 = pts_2d[i+1,1]

            x_data = np.array([x1+w[0], x1-w[0], x2-w[0], x2+w[0], x1+w[0]])
            y_data = np.array([y1+w[1], y1-w[1], y2-w[1], y2+w[1], y1+w[1]])

            if flip_xy:
                tmp = y_data
                y_data = x_data
                x_data = -tmp

            pp.plot(x_data, y_data, '-', alpha=alpha, color=color,
                    linewidth=linewidth)

if __name__ == '__main__':
    rospy.init_node('ode_arms_test')
    import hrl_common_code_darpa_m3.robot_config.six_link_planar as d_robot
    ode_arm = ODESimArm(d_robot)

    # test_jep = np.radians([30., 20., 45.])
    test_jep = np.radians(np.zeros(ode_arm.kinematics.n_jts))
    # test_jep = np.radians([90]*6)
    # test_jep = np.radians([0., 0., 0.])
    p, _ = ode_arm.kinematics.FK(test_jep, ode_arm.kinematics.n_jts)
    print 'p:', p.A1

    ode_arm.set_ep(test_jep)

max_stars_repo_name: gt-ros-pkg/hrl-haptic-manip

max_issues_repo_name: gt-ros-pkg/hrl-haptic-manip

max_forks_repo_name: gt-ros-pkg/hrl-haptic-manip


text:
# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for `alias.py`."""

from absl.testing import absltest
from absl.testing import parameterized

import chex
import jax
import jax.numpy as jnp

from optax._src import alias
from optax._src import numerics
from optax._src import update

def _setup_parabola(dtype):
  """Quadratic function as an optimization target."""
  initial_params = jnp.array([-1.0, 10.0, 1.0], dtype=dtype)
  final_params = jnp.array([1.0, -1.0, 1.0], dtype=dtype)

  if jnp.iscomplexobj(dtype):
    final_params *= 1 + 1j

  @jax.grad
  def get_updates(params):
    return jnp.sum(numerics.abs_sq(params - final_params))

  return initial_params, final_params, get_updates

def _setup_rosenbrock(dtype):
  """Rosenbrock function as an optimization target."""
  a = 1.0
  b = 100.0

  if jnp.iscomplexobj(dtype):
    a *= 1 + 1j

  initial_params = jnp.array([0.0, 0.0], dtype=dtype)
  final_params = jnp.array([a, a**2], dtype=dtype)

  @jax.grad
  def get_updates(params):
    return (numerics.abs_sq(a - params[0]) +
            b * numerics.abs_sq(params[1] - params[0]**2))

  return initial_params, final_params, get_updates

class AliasTest(chex.TestCase):

  @parameterized.product(
      (
          dict(opt_name='sgd', opt=lambda: alias.sgd(1e-3, 0.9)),
          dict(opt_name='adafactor', opt=lambda: alias.adafactor(5e-3)),
          dict(opt_name='adagrad', opt=lambda: alias.adagrad(1.0)),
          dict(opt_name='adam', opt=lambda: alias.adam(1e-1)),
          dict(opt_name='adamw', opt=lambda: alias.adamw(1e-1)),
          dict(opt_name='lars', opt=lambda: alias.lars(1.0)),
          dict(opt_name='lamb', opt=lambda: alias.lamb(1e-3)),
          dict(
              opt_name='noisy_sgd',
              opt=lambda: alias.noisy_sgd(1e-3, eta=1e-4)),
          dict(opt_name='rmsprop', opt=lambda: alias.rmsprop(5e-3)),
          dict(
              opt_name='rmsprop_momentum',
              opt=lambda: alias.rmsprop(5e-3, momentum=0.9)),
          dict(opt_name='fromage', opt=lambda: alias.fromage(5e-3)),
          dict(opt_name='adabelief', opt=lambda: alias.adabelief(1e-2)),
          dict(opt_name='radam', opt=lambda: alias.radam(5e-3)),
          dict(opt_name='sm3', opt=lambda: alias.sm3(1.0)),
          dict(opt_name='yogi', opt=lambda: alias.yogi(1e-1)),
          dict(
              opt_name='dpsgd',
              opt=lambda: alias.dpsgd(1e-3, 10.0, 0.001, 0, 0.2)),
      ),
      target=(_setup_parabola, _setup_rosenbrock),
      dtype=(jnp.float32, jnp.complex64),
  )
  def test_optimization(self, opt_name, opt, target, dtype):
    if (opt_name in ('fromage', 'noisy_sgd', 'sm3') and
        jnp.iscomplexobj(dtype)):
      raise absltest.SkipTest(
          f'{opt_name} does not support complex parameters.')

    opt = opt()
    initial_params, final_params, get_updates = target(dtype)

    @jax.jit
    def step(params, state):
      updates = get_updates(params)
      if opt_name == 'dpsgd':
        updates = updates[None]
      # Complex gradients need to be conjugated before being added to parameters
      # https://gist.github.com/wdphy16/118aef6fb5f82c49790d7678cf87da29
      updates = jax.tree_map(lambda x: x.conj(), updates)
      updates, state = opt.update(updates, state, params)
      params = update.apply_updates(params, updates)
      return params, state

    params = initial_params
    state = opt.init(params)
    for _ in range(10000):
      params, state = step(params, state)

    chex.assert_tree_all_close(params, final_params, rtol=3e-2, atol=3e-2)

  @parameterized.named_parameters([
      ('float32', 'float32'),
      ('bfloat16', 'bfloat16'),
      ('complex64', 'complex64'),
      ('None', None),
  ])
  def test_explicit_dtype(self, dtype):
    expected_dtype = jax.dtypes.canonicalize_dtype(dtype)  # None -> float32
    tx = alias.sgd(0.1, momentum=0.9, accumulator_dtype=dtype)
    trace_state, _ = tx.init(jnp.array([0.0, 0.0]))
    self.assertEqual(expected_dtype, trace_state.trace.dtype)
    tx = alias.adam(0.1, mu_dtype=dtype)
    adam_state, _ = tx.init(jnp.array([0.0, 0.0]))
    self.assertEqual(expected_dtype, adam_state.mu.dtype)
    tx = alias.adamw(0.1, mu_dtype=dtype)
    adam_state, _, _ = tx.init(jnp.array([0.0, 0.0]))
    self.assertEqual(expected_dtype, adam_state.mu.dtype)

if __name__ == '__main__':
  absltest.main()

max_stars_repo_name: eserie/optax

max_issues_repo_name: eserie/optax

max_forks_repo_name: eserie/optax

reason: import jax


text:
# -*- coding: utf-8 -*-
"""
Created on Wed Mar 23 12:53:10 2016

@author: DanielleT
"""

import matplotlib.pyplot as plt
import PIL
from PIL import Image
import os
import math
import pandas as pd
from pandas import *
import numpy as np
from sklearn.svm import SVC
from sklearn.preprocessing import MultiLabelBinarizer
import time
from sklearn import datasets
from sklearn.multiclass import OneVsRestClassifier
from sklearn.svm import LinearSVC

#LOAD DATA
print 'Loading Data....'
t = time.time()
#loads featurevectors
features1 = pd.read_csv('../../desktop/Features_data/caffe_features_train.csv',sep=',', header=None, iterator=True,chunksize=1000)
features =  concat(features1, ignore_index=True)
#loads links photo to business_id
train_photos = pd.read_csv('../../downloads/input/train_photo_to_biz_ids.csv',sep=',')
#loads business to labels
biz_ids = pd.read_csv('../../downloads/input/train.csv',sep=',')
#loads order of featurevectors
photo_order = pd.read_csv('../../desktop/Features_data/photo_order_train.csv',sep=',',header=None)
X_TRAIN = pd.merge(biz_ids, train_photos, on='business_id')
print time.time()-t,
print 'seconds needed.'

print 'Getting Business Averages....',
ti = time.time()
all_labels = [0] * 9
all_averages = [0] * 4096
t = 0
print 'Status: ',
for a in X_TRAIN.business_id.unique():          #for each business
    t = t+1
    data = X_TRAIN[(X_TRAIN.business_id == a)]
    labels = data['labels'].unique()
    if not isinstance(labels[0],str) and math.isnan(labels[0])==True:             #tocheckforbusinesses without labels
        continue
    labels = labels[0].split()
    labels = [int(x) for x in labels]
    labels = MultiLabelBinarizer().fit_transform([[0,1,2,3,4,5,6,7,8],labels])
    rev_labels = labels[1]
    all_labels = np.vstack((np.array(all_labels),np.array(rev_labels)))
    photos = data['photo_id']
    featavrg = [0] * 4096
    for photo in photos:
            photo_ref = str(photo) + ''.join('m.jpg')
            indx = photo_order.loc[photo_order[0] == photo_ref]
            if not(indx.empty):                     #photo deleted from collection
                ind = indx.index[0]
                feats = features.loc[ind,:]
                fts = feats.values.tolist()
                featavrg = np.vstack((np.array(featavrg),np.array(fts)))
    featavrg = featavrg[1:]
    averages = np.average(featavrg, axis=0)
    all_averages = np.vstack((np.array(all_averages),np.array(averages)))
    if t/100 == int(t/100):                  #prints multiples of 100
        print t,
print ti - time.time(),
print 'seconds needed.'

np.savetxt("labels_train_y.csv", all_labels[1:], delimiter=",")
np.savetxt("feats_train_x.csv", all_averages[1:], delimiter=",")

#toepassen SVM
print 'Training SVM....'
y = all_labels[1:]
x = all_averages[1:]
ti = time.time()
S = OneVsRestClassifier(LinearSVC(random_state=0)).fit(x, y)
score = S.score(x,y)
print score
print ti-time.time()
#
#
#
#
##-------------------------------PRODUCE TEST DATA
t = time.time()
#loads featurevectors
features1 = pd.read_csv('../../desktop/Features_data/caffe_features_test.csv',sep=',', header=None, iterator=True,chunksize=1000)
features_test =  concat(features1, ignore_index=True)
#loads links photo to business_id
test_photos = pd.read_csv('../../downloads/input/test_photo_to_biz.csv',sep=',')
#loads order of featurevectors
photo_order_test = pd.read_csv('../../desktop/Features_data/photo_order_test.csv',sep=',',header=None)
print t-time.time()

ti = time.time()
businesses = [0]
all_averages = [0] * 4096
rejected_photos = [0]
t = 0
for a in test_photos.business_id.unique():          #for each business
    t = t+1
    data = test_photos[(test_photos.business_id == a)]
    photos = data['photo_id']
    featavrg = [0] * 4096
    for photo in photos:
            photo_ref = str(photo) + ''.join('m.jpg')
            indx = photo_order_test.loc[photo_order_test[0] == photo_ref]
            if not(indx.empty):                     #photo deleted from collection
                ind = indx.index[0]
                feats = features_test.loc[ind,:]
                fts = feats.values.tolist()
                featavrg = np.vstack((np.array(featavrg),np.array(fts)))
            else:
                rejected_photos = np.vstack((np.array(rejected_photos),np.array(photo)))
    featavrg = featavrg[1:]
    averages = np.average(featavrg, axis=0)
    all_averages = np.vstack((np.array(all_averages),np.array(averages)))
    businesses = np.vstack((np.array(businesses),np.array(a)))
    print t
print ti - time.time()

#save averages
np.savetxt("biz_test.csv", businesses[1:], delimiter=",")
np.savetxt("rejected_photos_test.csv", rejected_photos[1:],delimiter = ",")
np.savetxt("feats_test_x.csv", all_averages[1:], delimiter=",")

##Perform
ti = time.time()
x = all_averages[1:]

predictions = S.predict(x)
print predictions[:5]
np.savetxt("predictions_test.csv", predictions, delimiter=",")
print time.time()-ti

#print outputfile
bizz = test_photos.business_id.unique()
a=0
arr = [0]*2
while True:
    strlabels = ""
    for i in range(0,8):
        pred = predictions[a]
        lab = pred[i]
        if lab==1:
            strlabels = strlabels + " " + str(i)
    currarr = [bizz[a],strlabels.lstrip()]
    arr = np.vstack((np.array(arr),np.array(currarr)))
    if bizz[a]==bizz[-1]:
        break
    a = a+1

arr[0]=['business_id','labels']

np.savetxt("subm.csv", arr, delimiter=",")

this = pd.DataFrame(data=arr[1:,:],columns=arr[0,:])  # 1st row as the column names
this.to_csv('feats_sub3.csv',index=False)

max_stars_repo_name: TheLaurens/Team-Brains

max_issues_repo_name: TheLaurens/Team-Brains

max_forks_repo_name: TheLaurens/Team-Brains


text:
""" Show deform grid and checkerboard of 2d registration Nellix data

"""

import os, time

import numpy as np
import visvis as vv
import pirt.reg # Python Image Registration Toolkit
from stentseg.utils.datahandling import select_dir, loadvol
import scipy
from scipy import ndimage
from lspeas.utils.vis_grid_deform import imagePooper

# Select the ssdf basedir
basedir = select_dir(r'E:\Nellix_chevas\CT_SSDF\SSDF_automated',
            r'D:\Nellix_chevas_BACKUP\CT_SSDF\SSDF_automated')

# Select dataset to register
ptcode = 'chevas_10'
ctcode, nr = '12months', 1
cropname = 'prox'

# Select dataset to register
what = 'phases'
phase1 = 3 # 30%
phase2 = 9 # 90%
slice = None #70
slicey = 107

# Load volumes
s = loadvol(basedir, ptcode, ctcode, cropname, what)

vols = []
phases = []
for key in dir(s):
    if key.startswith('vol'):
        print(key)
        # create vol with zoom in z-direction
        zscale = (s[key].sampling[0] / s[key].sampling[1]) # z / y
        # resample vol using spline interpolation, 3rd order piecewise polynomial
        vol_zoom = scipy.ndimage.interpolation.zoom(s[key],[zscale,1,1],'float32')
        s[key].sampling = [s[key].sampling[1],s[key].sampling[1],s[key].sampling[2]]

        # set scale and origin
        vol_zoom_type = vv.Aarray(vol_zoom, s[key].sampling, s[key].origin)

        # get slice 2d
        if slice:
            vol = vol_zoom_type[slice,:,:] # z y x
        elif slicey:
            vol = vol_zoom_type[:,slicey,:] # z y x
        else: # use mid slice z
            vol = vol_zoom_type[int(0.5*volOr.shape[0]),:,:]

        phases.append(key)
        vols.append(vol)

t0 = time.time()

# Initialize registration object
reg = pirt.reg.GravityRegistration(*vols)

reg.params.mass_transforms = 2  # 2nd order (Laplacian) triggers more at lines
reg.params.speed_factor = 1.0
reg.params.deform_wise = 'groupwise' # groupwise!
reg.params.mapping = 'backward'
reg.params.deform_limit = 1.0
reg.params.final_scale = 1.0  # We might set this a wee bit lower than 1 (but slower!)
reg.params.scale_sampling = 16
reg.params.final_grid_sampling = 20
reg.params.grid_sampling_factor = 0.5

# Go!
reg.register(verbose=1)

t1 = time.time()
print('Registration completed, which took %1.2f min.' % ((t1-t0)/60))

imagePooper(reg, fname='chevas_10_slicey107_grid15', gridsize=15, phase1=phase1, phase2=phase2)

max_stars_repo_name: almarklein/stentseg

max_issues_repo_name: almarklein/stentseg

max_forks_repo_name: almarklein/stentseg


text:
# -*- coding: UTF-8 -*-

"""
Store functions will take care of storage and post-processing related matters after a child model is trained.
"""

import shutil
import os
import numpy as np
import matplotlib.pyplot as plt
from ..plots import plot_hessian, plot_training_history
from .commonOps import unpack_data

def get_store_fn(arg):
    """The getter function that returns a callable store function from a string

    Parameters
    ----------
    arg : str
        The string identifier for a particular store function. Current choices are:
        - general
        - model_plot
        - minimal

    Returns
    -------
    callable
        A callable store function
    """
    if callable(arg) is True:
        return arg
    elif arg is None:
        return None
    elif arg.lower() == 'store_general' or arg.lower() == 'general':
        return store_general
    elif arg.lower() == 'store_regression' or arg.lower() == 'regression':
        return store_regression
    elif arg.lower() == 'store_with_hessian' or arg.lower() == 'hessian':
        return store_with_hessian
    elif arg.lower() == 'store_with_model_plot' or arg.lower() == 'model_plot':
        return store_with_model_plot
    elif arg.lower() == 'store_minimal' or arg.lower() == 'minimal':
        return store_minimal
    else:
        raise Exception('cannot understand store_fn: %s' % arg)

def store_with_model_plot(
        trial,
        model,
        hist,
        data,
        pred,
        loss_and_metrics,
        working_dir='.',
        save_full_model=False,
        *args, **kwargs
):
    par_dir = os.path.join(working_dir, 'weights', 'trial_%s' % trial)
    os.makedirs(par_dir, exist_ok=True)
    store_general(trial=trial,
                  model=model,
                  hist=hist,
                  data=data,
                  pred=pred,
                  loss_and_metrics=loss_and_metrics,
                  working_dir=working_dir,
                  save_full_model=save_full_model
                  )
    from tensorflow.keras.utils import plot_model
    plot_model(model, to_file=os.path.join(par_dir, "model_arc.png"), show_shapes=True, show_layer_names=True)

def store_with_hessian(
        trial,
        model,
        hist,
        data,
        pred,
        loss_and_metrics,
        working_dir='.',
        save_full_model=False,
        knowledge_func=None
):
    assert knowledge_func is not None, "`store_with_hessian` requires parsing the" \
                                       "knowledge function used."
    par_dir = os.path.join(working_dir, 'weights', 'trial_%s' % trial)
    store_general(trial=trial,
                  model=model,
                  hist=hist,
                  data=data,
                  pred=pred,
                  loss_and_metrics=loss_and_metrics,
                  working_dir=working_dir,
                  save_full_model=save_full_model
                  )
    # plot_training_history(hist, par_dir)
    from keras.utils import plot_model
    plot_model(model, to_file=os.path.join(par_dir, "model_arc.png"))

    plot_hessian(knowledge_func, os.path.join(par_dir, "hess.png"))
    knowledge_func._reset()
    plt.close('all')
    return

def store_general(
        trial,
        model,
        hist,
        data,
        pred,
        loss_and_metrics,
        working_dir='.',
        save_full_model=False,
        *args, **kwargs
):
    par_dir = os.path.join(working_dir, 'weights', 'trial_%s' % trial)
    if os.path.isdir(par_dir):
        shutil.rmtree(par_dir)
    os.makedirs(par_dir)
    if save_full_model:
        model.save(os.path.join(working_dir, 'weights', 'trial_%s' % trial, 'full_bestmodel.h5'))
    try:
        plot_training_history(hist, par_dir)
    except:
        # TODO: this still not working for non-keras models, e.g. EnasAnn/Cnn
        pass
    if os.path.isfile(os.path.join(working_dir, 'temp_network.h5')):
        shutil.move(os.path.join(working_dir, 'temp_network.h5'), os.path.join(par_dir, 'bestmodel.h5'))
    data = unpack_data(data, unroll_generator_y=True)
    metadata = data[2] if len(data) > 2 else None
    obs = data[1]
    write_pred_to_disk(
        os.path.join(par_dir, 'pred.txt'),
        pred, obs, metadata,
        loss_and_metrics
    )

def store_minimal(
        trial,
        model,
        working_dir='.',
        save_full_model=False,
        **kwargs
):
    par_dir = os.path.join(working_dir, 'weights', 'trial_%s' % trial)
    if os.path.isdir(par_dir):
        shutil.rmtree(par_dir)
    os.makedirs(par_dir)
    if save_full_model:
        model.save(os.path.join(working_dir, 'weights', 'trial_%s' % trial, 'full_bestmodel.h5'))
    if os.path.isfile(os.path.join(working_dir, 'temp_network.h5')):
        shutil.move(os.path.join(working_dir, 'temp_network.h5'), os.path.join(par_dir, 'bestmodel.h5'))

def write_pred_to_disk(fn, y_pred, y_obs, metadata=None, metrics=None):
    with open(fn, 'w') as f:
        if metrics is not None:
            f.write('\n'.join(['# {}: {}'.format(x, metrics[x]) for x in metrics]) + '\n')
        if type(y_pred) is list:
            y_pred = np.concatenate(y_pred, axis=1)
            y_obs = np.concatenate(y_obs, axis=1)
        if len(np.unique(y_obs)) < 10:  # is categorical
            str_format = "%i"
        else:
            str_format = "%.3f"

        f.write('pred\tobs\tmetadata\n')
        for i in range(len(y_pred)):
            if len(y_pred[i].shape) > 1 or y_pred[i].shape[0] > 1:
                y_pred_i = ','.join(['%.3f' % x for x in np.array(y_pred[i])])
                y_obs_i = ','.join([str_format % x for x in np.array(y_obs[i])])
            else:
                y_pred_i = '%.3f' % y_pred[i]
                y_obs_i = str_format % y_obs[i]
            if metadata:
                f.write('%s\t%s\t%s\n' % (y_pred_i, y_obs_i, metadata[i]))
            else:
                f.write('%s\t%s\t%s\n' % (y_pred_i, y_obs_i, 'NA'))