text:
from precise.skaters.portfoliostatic.equalport import equal_long_port
from precise.skaters.covarianceutil.covfunctions import cov_to_corrcoef
import numpy as np
import math
from precise.skaters.locationutil.vectorfunctions import normalize

def _buy_and_hold_port(port, j:int, y, s:dict, cov, port_kwargs):
    """

         Sporadically calls port() but uses buy and hold in between.
         For this to work, 'y' must be changes in log prices.

    :param port:
    :param j:
    :param y:
    :param s:
    :param cov:
    :param port_kwargs:
    :return:  w   Portfolio weights
    """
    if j==1:
        w = port(cov=cov, **port_kwargs)
        s = {}
        return w, s
    else:
        n_dim = len(y)
        if s.get('w') is None:
            s['multiplier']=[1 for _ in range(n_dim)]
            s['count']=0
            s['w'] = port(cov=cov, **port_kwargs)
            w = s['w']
            return w, s
        else:
            s['count'] = s['count']+1
            if s['count'] % j == 0:
                s['multiplier'] = [1 for _ in range(n_dim)]
                s['w'] = port(cov=cov, **port_kwargs)
                return s['w'], s
            else:
                s['multiplier'] = [ mi*math.exp(yi) for mi,yi in zip(s['multiplier'],y)]
                w = normalize( [ wi*mi for wi,mi in zip(s['w'],s['multiplier']) ] )
                return w, s

def static_cov_manager_factory_d0(y, s, f, port, e=1, f_kwargs:dict=None, port_kwargs:dict=None, n_cold=5, zeta=0.0, j=1):
    """
       Basic manager pattern ignoring mean.
       Expects to receive changes in log(price).
       If you have opinions on means, you'll have to incorporate them into variance somehow
       via the covariance skater.

          :param f     cov skater ("d0")
          :param port  portfolio constructor
          :param zeta  Compromise between cov and corr portfolio (zeta=0, use cov only)
          :param j     How often to run the static portfolio construction.
                       (Warning: if j>1 this will be nonsense unless 'y' represents changes in log prices)

       :returns w, s
    """

    if f_kwargs is None:
        f_kwargs = {}
    if port_kwargs is None:
        port_kwargs = {}
    if not s:
        s = {'f_state':{},
             'port_state':{},
             'count':0,
             'account_state':{}}

    x_mean, x_cov, s['f_state'] = f(y=y,s=s['f_state'], k=1, e=e, **f_kwargs)
    s['count']+=1
    if s['count']>=n_cold and (e>0):
        s_account = s['account_state']
        w, s_account = _buy_and_hold_port(port=port, y=y, j=j, s=s_account, cov=x_cov, port_kwargs=port_kwargs)
        s['account_state'] = s_account
        if zeta is not None and (zeta>0):
            x_diag = np.diag(x_cov)
            x_corr = cov_to_corrcoef(x_cov)
            w_corr_raw = port(cov=x_corr, **port_kwargs)
            w_corr = normalize(w_corr_raw/x_diag)
            w = (1-zeta)*w + zeta*w_corr
    else:
        w = equal_long_port(cov=x_cov)
    return w, s

max_stars_repo_name: microprediction/precise

max_issues_repo_name: microprediction/precise

max_forks_repo_name: microprediction/precise


text:
# -*- coding: utf-8 -*-
"""aula08_2021-05-03_segmentacao_grabcut.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1QZ8cyCKIftncqoD5eV3Vw7rDhc-7QZiD

# **Processamento de Imagens e Imagens**
Engenharia da Computação - 2021.01

## Segmentação com GrabCut

### Download das imagens de teste
"""

"""### Imports"""

import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

"""### Code"""

img1 = cv.imread('drive/MyDrive/Colab Notebooks/College/Signal-Image-Processing/image_dataset_publico/estatua.jpg')
img1 = cv.cvtColor(img1, cv.COLOR_BGR2RGB)

# Retangulo de Segmentacao
img2 = img1.copy()
p1 = (120, 100)
p2 = (350, 350)
img2 = cv.rectangle(img2, p1, p2, (255, 0, 0), 2)

# Parametros
mascara = np.zeros(img1.shape[:2], np.uint8)    # receber o resultado do algoritmo
bgModel = np.zeros((1, 65), np.float64)         # modelo background
fgModel = np.zeros((1, 65), np.float64)         # modelo do foreground
retangulo = p1 + p2

# Aplicar o Algoritmo GRABCUT
# cv.GC_INIT_WITH_RECT = retangulo de segmentacao
# cv.GC_INIT_WITH_MASK = pixels de segmentacao informados pelo usuario
cv.grabCut(img1, mascara, retangulo, bgModel, fgModel, 5, cv.GC_INIT_WITH_RECT)

# Separar o Background do Foreground
# Apos a execucao do algoritmo, a variavel mascara contém valores entre 0 e 3
# Os valores iguais a 0 (zero) ou 2 (dois) representam o Background
# Os valores iguais a 1 (um) ou 3 (três) representam o Foreground

# Filtrar os pixels de background
filtro = np.where( (mascara == 0) | (mascara == 2), 0, 1).astype('uint8')
img3 = img1.copy()
img3 = img3 * filtro[:, :, np.newaxis]

plt.figure(figsize=(10, 10))
plt.subplot(131), plt.imshow(img1)
plt.subplot(132), plt.imshow(img2)
plt.subplot(133), plt.imshow(img3)
plt.show()

"""### Exercises

#### Imports
"""

"""#### Code"""

ext: py

max_stars_repo_name: WesGtoX/Signal-Image-Processing

max_issues_repo_name: WesGtoX/Signal-Image-Processing

max_forks_repo_name: WesGtoX/Signal-Image-Processing

reason: import numpy


text:
#!/Users/ryan/anaconda2/envs/cdat_lite/bin/python

'''
Zelinka et al 2013 decomposition modified using oroginal code from Mark Zelinka

files required
cloud kernels from Zelinka et al 2012
control cosp run with ISCCP simulator
future cosp run with ISCCP simulator

clisccp variable dims: (240, 7, 7, 96, 192)
rsuscs/rsdscs variable dims: (240, 96, 192)
tas variable dims: (240, 96, 192)
output LW and SW feedback variable dims: (10, 90, 144)

RL - Sep 2017

'''

import cdms2 as cdms
import cdutil
import MV2 as MV
import numpy as np
import pylab as pl

# files
run = 'otb'
ker = 'cloud_kernels2.nc' # MZ kernel file
fctrl = run+'_ctrl_isccp.nc' # control ISCCP simulator output
fdble = run+'_dble_isccp.nc' # 2xCO2 ISCCP simulator output
direc="/Users/ryan/Desktop/2nd/cesm/data/cosp_data/"

ecs = 3.73949874
tau = [0.,0.3,1.3,3.6,9.4,23.,60.,380.]
CTP = [1000,800,680,560,440,310,180,50]
dlow = 5 # low cloud index, 3 for up to (including) 680hPa, 5 for 440hPa

# variables to load from files, last 2 are for surface clear sky albedo
vars = ["FISCCP1_COSP","TS","FSDSC","FSNSC"]

# helpful functions provided by MZ
###########################################################################
def add_cyclic(data):
    # Add Cyclic point around 360 degrees longitude:
    lons=data.getLongitude()[:]
    dx=np.gradient(lons)[-1]
    data2 = data(longitude=(0, dx+np.max(lons)), squeeze=True)
    return data2

###########################################################################
def nanarray(vector):

    # this generates a masked array with the size given by vector
    # example: vector = (90,144,28)

    # similar to this=NaN*ones(x,y,z) in matlab

    this=MV.zeros(vector)
    this=MV.masked_where(this==0,this)

    return this

###########################################################################
def map_SWkern_to_lon(Ksw,albcsmap):

    from scipy.interpolate import interp1d
    ## Map each location's clear-sky surface albedo to the correct albedo bin
    # Ksw is size 12,7,7,lats,3
    # albcsmap is size A,lats,lons
    albcs=np.arange(0.0,1.5,0.5)
    A=albcsmap.shape[0]
    TT=Ksw.shape[1]
    PP=Ksw.shape[2]
    lenlat=Ksw.shape[3]
    lenlon=albcsmap.shape[2]
    SWkernel_map=nanarray((A,TT,PP,lenlat,lenlon))
    for M in range(A):
        MM=M
        while MM>11:
            MM=MM-12
        for LA in range(lenlat):
            alon=albcsmap[M,LA,:]
            # interp1d can't handle mask but it can deal with NaN (?)
            try:
                alon2=MV.where(alon.mask,np.nan,alon)
            except:
                alon2=alon
            if np.ma.count(alon2)>1: # at least 1 unmasked value
                if len(pl.find(Ksw[MM,:,:,LA,:]>0))==0:
                    SWkernel_map[M,:,:,LA,:] = 0
                else:
                    f = interp1d(albcs,Ksw[MM,:,:,LA,:],axis=2,fill_value="extrapolate")
                    ynew = f(alon2.data)
                    ynew=MV.masked_where(alon2.mask,ynew)
                    SWkernel_map[M,:,:,LA,:] = ynew
            else:
                continue

    return SWkernel_map

def KT_decomposition_4D(c1,c2,Klw,Ksw):

    # this function takes in a (tau,CTP,lat,lon) matrix and performs the
    # decomposition of Zelinka et al 2013 doi:10.1175/JCLI-D-12-00555.1

    # reshape to be (CTP,tau,lat,lon)
    c1 = MV.transpose(c1,(1,0,2,3)) # control cloud fraction histogram
    c2 = MV.transpose(c2,(1,0,2,3)) # perturbed cloud fraction histogram
    Klw = MV.transpose(Klw,(1,0,2,3)) # LW Kernel histogram
    Ksw = MV.transpose(Ksw,(1,0,2,3)) # SW Kernel histogram

    P=c1.shape[0]
    T=c1.shape[1]

    c=c1
    sum_c=np.tile(MV.sum(MV.sum(c,1),0),(P,T,1,1))                                  # Eq. B2
    dc = c2-c1
    sum_dc=np.tile(MV.sum(MV.sum(dc,1),0),(P,T,1,1))
    dc_prop = c*(sum_dc/sum_c)
    dc_star = dc - dc_prop                                                          # Eq. B1

    # LW components
    Klw0 = np.tile(MV.sum(MV.sum(Klw*c/sum_c,1),0),(P,T,1,1))                       # Eq. B4
    Klw_prime = Klw - Klw0                                                          # Eq. B3
    this=MV.sum(Klw_prime*np.tile(MV.sum(c/sum_c,0),(P,1,1,1)),1)                   # Eq. B7a
    Klw_p_prime=np.tile(np.tile(this,(1,1,1,1)),(T,1,1,1))(order=[1,0,2,3])         # Eq. B7b
    that=np.tile(np.tile(MV.sum(c/sum_c,1),(1,1,1,1)),(T,1,1,1))(order=[1,0,2,3])   # Eq. B8a
    Klw_t_prime = np.tile(MV.sum(Klw_prime*that,0),(P,1,1,1))                       # Eq. B8b
    Klw_resid_prime = Klw_prime - Klw_p_prime - Klw_t_prime                         # Eq. B9
    dRlw_true = MV.sum(MV.sum(Klw*dc,1),0)                                          # LW total
    dRlw_prop = Klw0[0,0,:,:]*sum_dc[0,0,:,:]                                       # LW amount component
    dRlw_dctp = MV.sum(MV.sum(Klw_p_prime*dc_star,1),0)                             # LW altitude component
    dRlw_dtau = MV.sum(MV.sum(Klw_t_prime*dc_star,1),0)                             # LW optical depth component
    dRlw_resid = MV.sum(MV.sum(Klw_resid_prime*dc_star,1),0)                        # LW residual
    dRlw_sum = dRlw_prop + dRlw_dctp + dRlw_dtau + dRlw_resid                       # sum of LW components -- should equal LW total

    # SW components
    Ksw0 = np.tile(MV.sum(MV.sum(Ksw*c/sum_c,1),0),(P,T,1,1))                       # Eq. B4
    Ksw_prime = Ksw - Ksw0                                                          # Eq. B3
    this=MV.sum(Ksw_prime*np.tile(MV.sum(c/sum_c,0),(P,1,1,1)),1)                   # Eq. B7a
    Ksw_p_prime=np.tile(np.tile(this,(1,1,1,1)),(T,1,1,1))(order=[1,0,2,3])         # Eq. B7b
    that=np.tile(np.tile(MV.sum(c/sum_c,1),(1,1,1,1)),(T,1,1,1))(order=[1,0,2,3])   # Eq. B8a
    Ksw_t_prime = np.tile(MV.sum(Ksw_prime*that,0),(P,1,1,1))                       # Eq. B8b
    Ksw_resid_prime = Ksw_prime - Ksw_p_prime - Ksw_t_prime                         # Eq. B9
    dRsw_true = MV.sum(MV.sum(Ksw*dc,1),0)                                          # SW total
    dRsw_prop = Ksw0[0,0,:,:]*sum_dc[0,0,:,:]                                       # SW amount component
    dRsw_dctp = MV.sum(MV.sum(Ksw_p_prime*dc_star,1),0)                             # SW altitude component
    dRsw_dtau = MV.sum(MV.sum(Ksw_t_prime*dc_star,1),0)                             # SW optical depth component
    dRsw_resid = MV.sum(MV.sum(Ksw_resid_prime*dc_star,1),0)                        # SW residual
    dRsw_sum = dRsw_prop + dRsw_dctp + dRsw_dtau + dRsw_resid                       # sum of SW components -- should equal SW total

    dc_star = MV.transpose(dc_star,(1,0,2,3))
    dc_prop = MV.transpose(dc_prop,(1,0,2,3))

    return dRlw_true,dRlw_prop,dRlw_dctp,dRlw_dtau,dRlw_resid,dRsw_true,dRsw_prop,dRsw_dctp,dRsw_dtau,dRsw_resid

# Load in the Zelinka et al 2012 kernels:
f=cdms.open(direc+ker)
LWkernel=f('LWkernel')
SWkernel=f('SWkernel')
f.close()

# Load in variables for control simulation
fc = cdms.open(direc+fctrl,'r')

# Load in clisccp: ISCCP cloud area fraction (in atmosphere layer), surface air temperature (K)
clisccp1=fc(vars[0])
tas1 = fc(vars[1])

# Load in rsuscs, rsdscs
# rsuscs: surface upwelling shortwave flux in air assuming clear sky
# rsdscs: surface downwelling shortwave flux in air assuming clear sky
rsuscs1 = fc(vars[2])-fc(vars[3]) # net - down
rsdscs1 = fc(vars[3])
# "rsuscs/rsdscs variable dims:" (240, 96, 144)
albcs=np.arange(0.0,1.5,0.5) # the clear-sky albedos over which the kernel is computed

fc.close()

# Do the same for future
fd = cdms.open(direc+fdble,'r')
clisccp2=fd(vars[0])
tas2 = fd(vars[1])
fd.close()

# Define the cloud kernel axis attributes
lats=LWkernel.getLatitude()[:]
lons=np.arange(1.25,360,2.5)
grid = cdms.createGenericGrid(lats,lons)

# Make sure clisccp is in percent
sumclisccp1=MV.sum(MV.sum(clisccp1,2),1)
sumclisccp2=MV.sum(MV.sum(clisccp2,2),1)
if np.max(sumclisccp1) <= 1.:
    clisccp1 = clisccp1*100.
if np.max(sumclisccp2) <= 1.:
    clisccp2 = clisccp2*100.

# Compute climatological annual cycle:
avgclisccp1=cdutil.ANNUALCYCLE.climatology(clisccp1) #(12, TAU, CTP, LAT, LON)
avgclisccp2=cdutil.ANNUALCYCLE.climatology(clisccp2) #(12, TAU, CTP, LAT, LON)
del(clisccp1,clisccp2)

# Compute climatological annual cycle:
avgtas1=cdutil.ANNUALCYCLE.climatology(tas1) #(12, 90, 144)
avgtas2=cdutil.ANNUALCYCLE.climatology(tas2) #(12, 90, 144)
# Compute global annual mean tas anomalies
avgdtas = cdutil.averager(MV.average(avgtas2 - avgtas1,axis=0), axis='xy', weights='weighted') # (scalar)
print "globally averaged temperature anomalies: "+str(avgdtas)+" K"

# Compute clear-sky surface albedo
albcs1=rsuscs1/rsdscs1
avgalbcs1=cdutil.ANNUALCYCLE.climatology(albcs1) #(12, 90, 144)
avgalbcs1=MV.where(avgalbcs1>1.,1,avgalbcs1) # where(condition, x, y) is x where condition is true, y otherwise
avgalbcs1=MV.where(avgalbcs1<0.,0,avgalbcs1)
del(rsuscs1,rsdscs1,albcs1)

# Regrid everything to the kernel grid:
avgalbcs1 = add_cyclic(avgalbcs1)
avgclisccp1 = add_cyclic(avgclisccp1)
avgclisccp2 = add_cyclic(avgclisccp2)
avgalbcs1_grd = avgalbcs1.regrid(grid,regridTool="esmf",regridMethod = "linear")
clisccp1_r = avgclisccp1.regrid(grid,regridTool="esmf",regridMethod = "linear")
clisccp2_r = avgclisccp2.regrid(grid,regridTool="esmf",regridMethod = "linear")

# CESM cosp output default is ctp,tau,lat,lon - swap these axes to tau,ctp,lat,lon
clisccp1 = cdms.asVariable(np.transpose(clisccp1_r,(0,2,1,3,4)))
clisccp2 = cdms.asVariable(np.transpose(clisccp2_r,(0,2,1,3,4)))
clisccp1.setAxisList(clisccp1_r.getAxisList())

# LW kernel does not depend on albcs, just repeat the final dimension over longitudes:
LWK=np.tile(np.tile(LWkernel[:,:,:,:,0],(1,1,1,1,1)),(144,1,1,1,1))(order=[1,2,3,4,0])

# Use control albcs to map SW kernel to appropriate longitudes
SWK = map_SWkern_to_lon(SWkernel,avgalbcs1_grd)

print "Decomposing kernels..."

# decompose the kernels per month
all_fbks = np.zeros((12,10,clisccp1.shape[-2],144))
all_fbks_low = np.zeros((12,10,clisccp1.shape[-2],144)) #low cloud feedback
all_fbks_nonlow = np.zeros_like(all_fbks_low) # non low cloud feedback
for i in range(12):
    # decompose low clouds by high cloud top pressures
    decomp_fbks = KT_decomposition_4D(clisccp1[i,:,:dlow,:,:],clisccp2[i,:,:dlow,:,:],
                                              LWK[i,:,:dlow,:,:],SWK[i,:,:dlow,:,:])
    var_out = np.array(decomp_fbks)
    all_fbks_low[i,:,:,:] = var_out[:,:,:]
    # decompose the non low clouds
    decomp_fbks = KT_decomposition_4D(clisccp1[i,:,dlow:,:,:],clisccp2[i,:,dlow:,:,:],
                                              LWK[i,:,dlow:,:,:],SWK[i,:,dlow:,:,:])
    var_out = np.array(decomp_fbks)
    all_fbks_nonlow[i,:,:,:] = var_out[:,:,:]
    decomp_fbks = KT_decomposition_4D(clisccp1[i,:,:,:,:],clisccp2[i,:,:,:,:],
                                              LWK[i,:,:,:,:],SWK[i,:,:,:,:])
    var_out = np.array(decomp_fbks)
    all_fbks[i,:,:,:] = var_out[:,:,:]

print "done"
var_out1 = np.average(all_fbks_low,axis=0)
var_out1 /= ecs
var_out2 = np.average(all_fbks_nonlow,axis=0)
var_out2 /= ecs
var_out3 = np.average(all_fbks,axis=0)
var_out3 /= ecs

print "Saving to netCDF file"
from netCDF4 import Dataset

# Write to a netcdf file
dataset = Dataset(run+'_cloud_fbk_decomp.nc','w')

# Create dimensions
i = dataset.createDimension('i',var_out.shape[0])
lat = dataset.createDimension('lat',len(lats))
lon = dataset.createDimension('lon',len(lons))

# Create and Write dimensional variables
LAT = dataset.createVariable('LAT', np.float32,('lat',))
LON = dataset.createVariable('LON', np.float32,('lon',))
LAT[:] = lats
LON[:] = lons

FB_LOW = dataset.createVariable('FB_LOW',np.float32,('i','lat','lon'))
FB_LOW.units = 'W/m2/K'
FB_LOW.description = "Low cloud feedback decomposition from KT_decomposition_4D function.\nDimensions are [feedback type,lat,lon]"
FB_LOW[:,:,:] = var_out1

FB_NLOW = dataset.createVariable('FB_NLOW',np.float32,('i','lat','lon'))
FB_NLOW.units = 'W/m2/K'
FB_NLOW.description = "Non-low cloud feedback decomposition from KT_decomposition_4D function.\nDimensions are [feedback type,lat,lon]"
FB_NLOW[:,:,:] = var_out2

FB_ALL = dataset.createVariable('FB_ALL',np.float32,('i','lat','lon'))
FB_ALL.units = 'W/m2/K'
FB_ALL.description = "Total cloud feedback decomposition from KT_decomposition_4D function.\nDimensions are [feedback type,lat,lon]"
FB_ALL[:,:,:] = var_out3

dataset.close()

max_stars_repo_name: stefan-hofer/cloud-radiative-kernels

max_issues_repo_name: stefan-hofer/cloud-radiative-kernels

max_forks_repo_name: stefan-hofer/cloud-radiative-kernels

reason: import numpy,from scipy


text:
'''
This script augments the dataset by pitch-shifting the individual singers recordings 2 semitones above and below the original
pitch. It uses the MUDA python package for music data augmentation.

This is a stand-alone script that is executed even before the setup process.

'''

import muda
import jams
import numpy as np
import pandas as pd

import os
import argparse
import csv

import utils

def create_jams(times, freqs, outfile):

    track_duration = times[-1] + (times[-1] - times[-2])

    jam = jams.JAMS()
    jam.file_metadata.duration = track_duration

    pitch_a = jams.Annotation(namespace='pitch_contour')
    pitch_a.annotation_metadata.data_source = "Tony pitch estimation + manual correction"
    pitch_a.annotation_metadata.annotation_tools = "Tony"

    for t, p in zip(times, freqs):

        if p != 0:
            pitch_a.append(
                time=t,
                duration=0.0,
                value={'index': 0, 'frequency': p, 'voiced': True},
                confidence=1.0
            )
        else:
            pitch_a.append(
                time=t,
                duration=0.0,
                value={'index': 0, 'frequency': p, 'voiced': False},
                confidence=1.0
            )

    jam.annotations.append(pitch_a)

    jam.save(outfile)

#####################################################################################

def read_annotations_f0(annot_fname, annot_path, dataset=None):

    if annot_fname.endswith('f0'):

        if dataset == 'ECS':
            # loadtxt fails with some ECS files
            # annotation = np.loadtxt(os.path.join(annot_path, annot_fname))
            annotation = []
            with open(os.path.join(annot_path, annot_fname), newline='\n') as f:
                reader = csv.reader(f, delimiter='\t')
                for line in reader:
                    annotation.append([float(line[0]), float(line[1])])
                annotation = np.array(annotation)
        else:
            print("CSD: {}".format(annot_path))
            annotation = np.loadtxt(os.path.join(annot_path, annot_fname))

    elif annot_fname.endswith('csv'):
        annotation = pd.read_csv(os.path.join(annot_path, annot_fname), header=None).values
    else:
        print("Invalid annotation file format for {}".format(annot_fname))

    return annotation[:, 0], annotation[:, 1]

#####################################################################################

def pitch_shifting(audio_fname, jams_fname, audio_folder, jams_folder, n_samples=5, l=-2, u=2):

    print(audio_folder, audio_fname)

    orig = muda.load_jam_audio(os.path.join(jams_folder, jams_fname),
                               os.path.join(audio_folder, audio_fname)
                               )

    pitchshift = muda.deformers.LinearPitchShift(n_samples=n_samples, lower=l, upper=u)

    for i, jam_out in enumerate(pitchshift.transform(orig)):

        muda.save(os.path.join(audio_folder, '{}_{}'.format(i, audio_fname)),
                  os.path.join(jams_folder, '{}_{}'.format(i, jams_fname)),
                  jam_out)

#####################################################################################

def add_unvoiced_frames(annot_path, audio_path, format='csv', fs=22050.0):

    if not os.path.exists(os.path.join(annot_path, 'constant_timebase')):
        os.mkdir(os.path.join(annot_path, 'constant_timebase'))

    for fname in os.listdir(annot_path):
        if not fname.endswith(format): continue
        if not 'smoothedpitchtrack' in fname: continue

        utils.pyin_to_unvoiced(annot_path, fname, audio_path, fname.replace('_vamp_pyin_pyin_smoothedpitchtrack.csv',
                                                                            '.wav'), fs=fs)

#####################################################################################

def main(args):

    # fix pYIN annotations for BC and BSQ and then keep going. After this, annot folder becomes constant_timebase
    if args.dataset == 'BC' or args.dataset == 'BSQ':

        add_unvoiced_frames(args.path_to_annotations, args.path_to_audio, fs=44100.0)

        # once unvoiced frames are fixed, data augmentation steps for audio and annotations (CSV)
        for fn in os.listdir(os.path.join(args.path_to_annotations, 'constant_timebase')):

            if not fn.endswith('csv'): continue

            orig_times, orig_freqs = read_annotations_f0(fn, os.path.join(args.path_to_annotations, 'constant_timebase'))

            new_fname = fn.replace('_vamp_pyin_pyin_smoothedpitchtrack.csv', '_pyin.jams')
            outfile = os.path.join(args.path_to_annotations, 'constant_timebase', new_fname)
            create_jams(orig_times, orig_freqs, outfile)

            # step 3 is pitch-shifting audio and annotations accordingly
            pitch_shifting(new_fname.replace('_pyin.jams', '.wav'), new_fname, args.path_to_audio,
                           os.path.join(args.path_to_annotations, 'constant_timebase'))

    elif args.dataset == 'ECS':
        # step 2 is converting annotation files to jams
        for fn in os.listdir(args.path_to_annotations):

            if not fn.endswith('f0'): continue

            orig_times, orig_freqs = read_annotations_f0(fn, args.path_to_annotations, dataset='ECS')

            outfile = os.path.join(args.path_to_annotations, fn.replace('f0', 'jams'))
            create_jams(orig_times, orig_freqs,  outfile)

            # step 3 is pitch-shifting audio and annotations accordingly

            pitch_shifting(fn.replace('f0', 'wav'), fn.replace('f0', 'jams'), args.path_to_audio,
                           args.path_to_annotations)

    else:
        # step 2 is converting annotation files to jams
        for fn in os.listdir(args.path_to_annotations):

            if not fn.endswith('f0'):
                if not fn.endswith('csv'): continue

            orig_times, orig_freqs = read_annotations_f0(fn, args.path_to_annotations)

            if fn.endswith('f0'):
                outfile = os.path.join(args.path_to_annotations, fn.replace('f0', 'jams'))
                create_jams(orig_times, orig_freqs,  outfile)

            elif fn.endswith('csv'):
                outfile = os.path.join(args.path_to_annotations, fn.replace('csv', 'jams'))
                create_jams(orig_times, orig_freqs, outfile)

            # step 3 is pitch-shifting audio and annotations accordingly
            if fn.endswith('f0'):
                pitch_shifting(fn.replace('f0', 'wav'), fn.replace('f0', 'jams'), args.path_to_audio,
                               args.path_to_annotations)
            elif fn.endswith('csv'):
                pitch_shifting(fn.replace('csv', 'wav'), fn.replace('csv', 'jams'), args.path_to_audio,
                               args.path_to_annotations)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Two steps: (1) Convert f0 annotations to jams format for further use, (2) pitch-shifting for data augmentation.")

    parser.add_argument("--f0-path",
                        dest='path_to_annotations',
                        type=str,
                        help="Path to folder with f0 files. ")

    parser.add_argument("--audio-path",
                        dest='path_to_audio',
                        type=str,
                        help="Path to folder with audio files. ")

    parser.add_argument("--dataset",
                        dest="dataset",
                        type=str,
                        help="Indicate the dataset to process: BC / BSQ / CSD / ECS / DCS")
    '''
    parser.add_argument("--pyin",
                        dest='pyin',
                        type=str,
                        help="If F0-trajectories come from pYIN they might not have unvoiced frames as 0. If 'yes', the code takes care of this. If 'no', "
                             "it assumes unvoiced frames are OK and skips this.")
    '''

    main(parser.parse_args())


text:
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"

import pandas as pd
import numpy as np
from torch.utils.data import DataLoader
from data_loading import create_datasets, SeriesDataset
from config import get_config
from trainer import ESRNNTrainer
from model import ESRNN
from DRNN import DRNN
import time

import os
import time
import numpy as np
import copy
import torch
import torch.nn as nn
from loss_modules import PinballLoss, sMAPE, np_sMAPE
from utils.logger import Logger
import pandas as pd

print('loading config')
config = get_config('Daily')

config['output_size'] = 30

print('loading data')
info = pd.read_csv('data/energy_daily.csv')

train_path = 'data/train/%s-train.csv' % (config['variable'])
test_path = 'data/test/%s-test.csv' % (config['variable'])

df = pd.read_csv('data/energy_daily.csv')

df = df['W']

df.max()

df.min()

# +
#df = (df - df.min())/(df.max() - df.min()) #model normalizes
# -

values = df.values

# data = []
# seq_len = 30
# for i in range(len(values)-seq_len+1):
#     #print(list(values[i:i+seq_len]))
#     data.append(list(values[i:i+seq_len].reshape(-1)))

# df = pd.DataFrame(data)

# df.head()

train_df = pd.DataFrame(df[:int(0.8*len(df))])
test_df = pd.DataFrame(df[int(0.8*len(df)):])

train_df.to_csv(train_path, index = False)
test_df.to_csv(test_path, index = False)

train = train_df.values.T
test = test_df.values.T

train.shape

len(train)

test.shape

val = train[:, -config['output_size']:]
train = train[:, :-config['output_size']]

train.shape

val.shape

# train, val, test = create_datasets(train_path, test_path, config['output_size'])

import torch
from torch.utils.data import Dataset

class SeriesDataset(Dataset):

    def __init__(self, dataTrain, dataVal, dataTest, variable, chop_value, device):
        #dataTrain, mask = chop_series(dataTrain, chop_value)

        #self.dataInfoCatOHE = pd.get_dummies(info[info['SP'] == variable]['category'])
        #self.dataInfoCatHeaders = np.array([i for i in self.dataInfoCatOHE.columns.values])
        #self.dataInfoCat = torch.from_numpy(self.dataInfoCatOHE[mask].values).float()
        self.dataTrain = [torch.tensor(dataTrain[i]) for i in range(len(dataTrain))]  # ALREADY MASKED IN CHOP FUNCTION
        self.dataVal = [torch.tensor(dataVal[i]) for i in range(len(dataVal))]
        self.dataTest = [torch.tensor(dataTest[i]) for i in range(len(dataTest))]
        self.device = device

    def __len__(self):
        return len(self.dataTrain)

    def __getitem__(self, idx):
        return self.dataTrain[idx].to(self.device), \
                self.dataVal[idx].to(self.device), \
                self.dataTest[idx].to(self.device), \
                idx
                #self.dataInfoCat[idx].to(self.device), \

dataset = SeriesDataset(train, val, test, config['variable'], config['chop_val'], config['device'])
dataloader = DataLoader(dataset, batch_size=config['batch_size'], shuffle=True)

for batch in dataloader:
    print(batch[0].shape)

len(dataset)

run_id = str(int(time.time())) #UNIX time
run_id

config['dilations']

for grp_num in range(len(config['dilations'])):
    print(grp_num)

class ResidualDRNN(nn.Module):
    def __init__(self, config):
        super(ResidualDRNN, self).__init__()
        self.config = config

        layers = []
        for grp_num in range(len(self.config['dilations'])):

            if grp_num == 0:
                input_size = self.config['input_size']# + self.config['num_of_categories']
            else:
                input_size = self.config['state_hsize']

            #input_size = self.config['input_size']

            l = DRNN(input_size,
                     self.config['state_hsize'],
                     n_layers=len(self.config['dilations'][grp_num]),
                     dilations=self.config['dilations'][grp_num],
                     cell_type=self.config['rnn_cell_type'])

            layers.append(l)

        self.rnn_stack = nn.Sequential(*layers)

    def forward(self, input_data):
        for layer_num in range(len(self.rnn_stack)):
            residual = input_data
            out, _ = self.rnn_stack[layer_num](input_data)
            if layer_num > 0:
                out += residual
            input_data = out
        return out

class ESRNN(nn.Module):
    def __init__(self, num_series, config):
        super(ESRNN, self).__init__()
        self.config = config
        self.num_series = num_series
        self.add_nl_layer = self.config['add_nl_layer']

        init_lev_sms = []
        init_seas_sms = []
        init_seasonalities = []

        for i in range(num_series):
            init_lev_sms.append(nn.Parameter(torch.Tensor([0.5]), requires_grad=True))
            init_seas_sms.append(nn.Parameter(torch.Tensor([0.5]), requires_grad=True))
            init_seasonalities.append(nn.Parameter((torch.ones(config['seasonality']) * 0.5), requires_grad=True))

        self.init_lev_sms = nn.ParameterList(init_lev_sms)
        self.init_seas_sms = nn.ParameterList(init_seas_sms)
        self.init_seasonalities = nn.ParameterList(init_seasonalities)

        self.nl_layer = nn.Linear(config['state_hsize'],
                                  config['state_hsize'])
        self.act = nn.Tanh()
        self.scoring = nn.Linear(config['state_hsize'], config['output_size'])

        self.logistic = nn.Sigmoid()

        self.resid_drnn = ResidualDRNN(self.config)

    def forward(self, train, val, test, idxs, testing=False):
        # GET THE PER SERIES PARAMETERS
        lev_sms = self.logistic(torch.stack([self.init_lev_sms[idx] for idx in idxs]).squeeze(1))
        seas_sms = self.logistic(torch.stack([self.init_seas_sms[idx] for idx in idxs]).squeeze(1))
        init_seasonalities = torch.stack([self.init_seasonalities[idx] for idx in idxs])

        seasonalities = []
        # PRIME SEASONALITY
        for i in range(self.config['seasonality']):
            seasonalities.append(torch.exp(init_seasonalities[:, i]))
        seasonalities.append(torch.exp(init_seasonalities[:, 0]))

        if testing:
            train = torch.cat((train, val), dim=1)

        train = train.float()

        levs = []
        log_diff_of_levels = []

        levs.append(train[:, 0] / seasonalities[0])
        for i in range(1, train.shape[1]):
            # CALCULATE LEVEL FOR CURRENT TIMESTEP TO NORMALIZE RNN
            new_lev = lev_sms * (train[:, i] / seasonalities[i]) + (1 - lev_sms) * levs[i - 1]
            levs.append(new_lev)

            # STORE DIFFERENCE TO PENALIZE LATER
            log_diff_of_levels.append(torch.log(new_lev / levs[i - 1]))

            # CALCULATE SEASONALITY TO DESEASONALIZE THE DATA FOR RNN
            seasonalities.append(seas_sms * (train[:, i] / new_lev) + (1 - seas_sms) * seasonalities[i])

        seasonalities_stacked = torch.stack(seasonalities).transpose(1, 0)
        levs_stacked = torch.stack(levs).transpose(1, 0)

        loss_mean_sq_log_diff_level = 0
        if self.config['level_variability_penalty'] > 0:
            sq_log_diff = torch.stack(
                [(log_diff_of_levels[i] - log_diff_of_levels[i - 1]) ** 2 for i in range(1, len(log_diff_of_levels))])
            loss_mean_sq_log_diff_level = torch.mean(sq_log_diff)

        if self.config['output_size'] > self.config['seasonality']:
            start_seasonality_ext = seasonalities_stacked.shape[1] - self.config['seasonality']
            end_seasonality_ext = start_seasonality_ext + self.config['output_size'] - self.config['seasonality']
            seasonalities_stacked = torch.cat((seasonalities_stacked, seasonalities_stacked[:, start_seasonality_ext:end_seasonality_ext]),
                                              dim=1)

        window_input_list = []
        window_output_list = []
        for i in range(self.config['input_size'] - 1, train.shape[1]):
            input_window_start = i + 1 - self.config['input_size']
            input_window_end = i + 1

            train_deseas_window_input = train[:, input_window_start:input_window_end] / seasonalities_stacked[:,
                                                                                input_window_start:input_window_end]
            train_deseas_norm_window_input = (train_deseas_window_input / levs_stacked[:, i].unsqueeze(1))
            #train_deseas_norm_cat_window_input = torch.cat((train_deseas_norm_window_input), dim=1)
            window_input_list.append(train_deseas_norm_window_input)

            output_window_start = i + 1
            output_window_end = i + 1 + self.config['output_size']

            if i < train.shape[1] - self.config['output_size']:
                train_deseas_window_output = train[:, output_window_start:output_window_end] / \
                                             seasonalities_stacked[:, output_window_start:output_window_end]
                train_deseas_norm_window_output = (train_deseas_window_output / levs_stacked[:, i].unsqueeze(1))
                window_output_list.append(train_deseas_norm_window_output)

        window_input = torch.cat([i.unsqueeze(0) for i in window_input_list], dim=0)
        window_output = torch.cat([i.unsqueeze(0) for i in window_output_list], dim=0)

        self.train()
        network_pred = self.series_forward(window_input[:-self.config['output_size']])
        network_act = window_output

        self.eval()
        network_output_non_train = self.series_forward(window_input)

        # USE THE LAST VALUE OF THE NETWORK OUTPUT TO COMPUTE THE HOLDOUT PREDICTIONS
        hold_out_output_reseas = network_output_non_train[-1] * seasonalities_stacked[:, -self.config['output_size']:]
        hold_out_output_renorm = hold_out_output_reseas * levs_stacked[:, -1].unsqueeze(1)

        hold_out_pred = hold_out_output_renorm * torch.gt(hold_out_output_renorm, 0).float()
        hold_out_act = test if testing else val

        hold_out_act_deseas = hold_out_act.float() / seasonalities_stacked[:, -self.config['output_size']:]
        hold_out_act_deseas_norm = hold_out_act_deseas / levs_stacked[:, -1].unsqueeze(1)

        self.train()
        # RETURN JUST THE TRAINING INPUT RATHER THAN THE ENTIRE SET BECAUSE THE HOLDOUT IS BEING GENERATED WITH THE REST
        return network_pred, \
               network_act, \
               (hold_out_pred, network_output_non_train), \
               (hold_out_act, hold_out_act_deseas_norm), \
               loss_mean_sq_log_diff_level

    def series_forward(self, data):
        data = self.resid_drnn(data)
        if self.add_nl_layer:
            data = self.nl_layer(data)
            data = self.act(data)
        data = self.scoring(data)
        return data

num_series = 1

# +
init_lev_sms = []
init_seas_sms = []
init_seasonalities = []

for i in range(num_series):
    init_lev_sms.append(nn.Parameter(torch.Tensor([0.5]), requires_grad=True))
    init_seas_sms.append(nn.Parameter(torch.Tensor([0.5]), requires_grad=True))
    init_seasonalities.append(nn.Parameter((torch.ones(config['seasonality']) * 0.5), requires_grad=True))
# -

init_lev_sms
init_seas_sms
init_seasonalities

init_lev_sms = nn.ParameterList(init_lev_sms)
init_seas_sms = nn.ParameterList(init_seas_sms)
init_seasonalities = nn.ParameterList(init_seasonalities)

nl_layer = nn.Linear(config['state_hsize'], config['state_hsize'])
act = nn.Tanh()
scoring = nn.Linear(config['state_hsize'], config['output_size'])
logistic = nn.Sigmoid()
resid_drnn = ResidualDRNN(config)

for train, val, test, idxs in dataloader:
    print(idxs)

# GET THE PER SERIES PARAMETERS
lev_sms = logistic(torch.stack([init_lev_sms[idx] for idx in idxs]).squeeze(1))
seas_sms = logistic(torch.stack([init_seas_sms[idx] for idx in idxs]).squeeze(1))
init_seasonalities = torch.stack([init_seasonalities[idx] for idx in idxs])

lev_sms
seas_sms
init_seasonalities

seasonalities = []
# PRIME SEASONALITY
for i in range(config['seasonality']):
    seasonalities.append(torch.exp(init_seasonalities[:, i]))
seasonalities.append(torch.exp(init_seasonalities[:, 0]))
seasonalities

train = train.float()
train.shape

testing = False

if testing:
    train = torch.cat((train, val), dim=1)
train.shape

train[:, 0]

levs = []
log_diff_of_levels = []
levs.append(train[:, 0] / seasonalities[0])

train[:, 0]
seasonalities[0]
train[:, 0] / seasonalities[0]

levs

lev_sms # = alpha

seas_sms # = gamma

for i in range(1, train.shape[1]):
    # CALCULATE LEVEL FOR CURRENT TIMESTEP TO NORMALIZE RNN
    new_lev = lev_sms * (train[:, i] / seasonalities[i]) + (1 - lev_sms) * levs[i - 1] #level update equation
    levs.append(new_lev)

    # STORE DIFFERENCE TO PENALIZE LATER
    log_diff_of_levels.append(torch.log(new_lev / levs[i - 1]))

    # CALCULATE SEASONALITY TO DESEASONALIZE THE DATA FOR RNN
    seasonalities.append(seas_sms * (train[:, i] / new_lev) + (1 - seas_sms) * seasonalities[i]) #seas update eqn

len(seasonalities) # for each step of the input, we use the seas of 'seasonality' steps back (in our case 7) 485=7+478

seasonalities

seasonalities_stacked = torch.stack(seasonalities).transpose(1, 0) #torch.stack stacks the elements of a list to
                                                                   #form a tensor

levs_stacked = torch.stack(levs).transpose(1, 0)

loss_mean_sq_log_diff_level = 0
if config['level_variability_penalty'] > 0:
    sq_log_diff = torch.stack(
        [(log_diff_of_levels[i] - log_diff_of_levels[i - 1]) ** 2 for i in range(1, len(log_diff_of_levels))])
    loss_mean_sq_log_diff_level = torch.mean(sq_log_diff)

loss_mean_sq_log_diff_level

config['output_size']
config['seasonality']

if config['output_size'] > config['seasonality']:
    start_seasonality_ext = seasonalities_stacked.shape[1] - config['seasonality']
    end_seasonality_ext = start_seasonality_ext + config['output_size'] - config['seasonality']
    seasonalities_stacked = torch.cat((seasonalities_stacked, seasonalities_stacked[:, start_seasonality_ext:end_seasonality_ext]),
                                      dim=1)

start_seasonality_ext

end_seasonality_ext

seasonalities_stacked[:].shape #added more length to make longer predictions

window_input_list = []
window_output_list = []
for i in range(config['input_size'] - 1, train.shape[1]):
    input_window_start = i + 1 - config['input_size']
    input_window_end = i + 1

    train_deseas_window_input = train[:, input_window_start:input_window_end] / seasonalities_stacked[:,
                                                                                input_window_start:input_window_end]
    train_deseas_norm_window_input = (train_deseas_window_input / levs_stacked[:, i].unsqueeze(1))
    #train_deseas_norm_cat_window_input = torch.cat((train_deseas_norm_window_input), dim=1)
    window_input_list.append(train_deseas_norm_window_input)

    output_window_start = i + 1
    output_window_end = i + 1 + config['output_size']

    if i < train.shape[1] - config['output_size']:
        train_deseas_window_output = train[:, output_window_start:output_window_end] / \
                                     seasonalities_stacked[:, output_window_start:output_window_end]
        train_deseas_norm_window_output = (train_deseas_window_output / levs_stacked[:, i].unsqueeze(1))
        window_output_list.append(train_deseas_norm_window_output)

len(window_input_list) #rolling window works correctly

len(window_output_list)

window_input = torch.cat([i.unsqueeze(0) for i in window_input_list], dim=0)
window_output = torch.cat([i.unsqueeze(0) for i in window_output_list], dim=0)

window_input.shape #now I get where we got 472 and 7

window_output.shape

add_nl_layer = config['add_nl_layer']

def series_forward(data): #series_forward calls resid_rnn and adds nl_layer
    data = resid_drnn(data)
    if add_nl_layer:
        data = nl_layer(data)
        data = act(data)
    data = scoring(data) #scoring converts from hidden_dim shape to output_dim shape
    return data

#self.train() since ESRNN inherits nn.Module it has train and eval funcs
network_pred = series_forward(window_input[:-config['output_size']])

window_input[:-config['output_size']].shape

network_pred.shape

network_act = window_output #true values
network_act.shape

window_input.shape

#self.eval()
network_output_non_train = series_forward(window_input)

network_pred.shape

network_output_non_train.shape

window_output.shape

network_act.shape

# USE THE LAST VALUE OF THE NETWORK OUTPUT TO COMPUTE THE HOLDOUT PREDICTIONS
# reseasonalize and renormalize by multiplying by seas and lev
hold_out_output_reseas = network_output_non_train[-1] * seasonalities_stacked[:, -config['output_size']:]
hold_out_output_renorm = hold_out_output_reseas * levs_stacked[:, -1].unsqueeze(1)

network_output_non_train.shape

hold_out_output_renorm

torch.gt(hold_out_output_renorm, 0)

hold_out_pred = hold_out_output_renorm * torch.gt(hold_out_output_renorm, 0).float()
hold_out_act = test if testing else val

hold_out_act

hold_out_pred

hold_out_act_deseas = hold_out_act.float() / seasonalities_stacked[:, -config['output_size']:]
hold_out_act_deseas_norm = hold_out_act_deseas / levs_stacked[:, -1].unsqueeze(1)

self.train()

        self.train()
        # RETURN JUST THE TRAINING INPUT RATHER THAN THE ENTIRE SET BECAUSE THE HOLDOUT IS BEING GENERATED WITH THE REST
        return network_pred, \
               network_act, \
               (hold_out_pred, network_output_non_train), \
               (hold_out_act, hold_out_act_deseas_norm), \
               loss_mean_sq_log_diff_level

a = torch.randn(2,5,6)
a
a.mean(dim = 0)

class ESRNNTrainer(nn.Module):
    def __init__(self, model, dataloader, run_id, config):
        super(ESRNNTrainer, self).__init__()
        self.model = model.to(config['device'])
        self.config = config
        self.dl = dataloader
        #self.ohe_headers = ohe_headers
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=config['learning_rate'])
        # self.optimizer = torch.optim.ASGD(self.model.parameters(), lr=config['learning_rate'])
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer,
                                                         step_size=config['lr_anneal_step'],
                                                         gamma=config['lr_anneal_rate'])
        self.criterion = PinballLoss(self.config['training_tau'],
                                     self.config['output_size'] * self.config['batch_size'], self.config['device'])
        self.epochs = 0
        self.max_epochs = config['num_of_train_epochs']
        self.run_id = str(run_id)
        self.prod_str = 'prod' if config['prod'] else 'dev'
        self.log = Logger("../logs/train%s%s%s" % (self.config['variable'], self.prod_str, self.run_id))
        self.csv_save_path = None

    def train_epochs(self):
        max_loss = 1e8
        start_time = time.time()
        for e in range(self.max_epochs):
            self.scheduler.step()

            epoch_loss = self.train()
            if epoch_loss < max_loss:
                self.save()
            epoch_val_loss = self.val()
            if e == 0:
                file_path = os.path.join(self.csv_save_path, 'validation_losses.csv')
                with open(file_path, 'w') as f:
                    f.write('epoch,training_loss,validation_loss\n')
            with open(file_path, 'a') as f:
                f.write(','.join([str(e), str(epoch_loss), str(epoch_val_loss)]) + '\n')
        print('Total Training Mins: %5.2f' % ((time.time()-start_time)/60))

    def train(self):
        self.model.train()
        epoch_loss = 0
        for batch_num, (train, val, test, idx) in enumerate(self.dl):
            start = time.time()
            print("Train_batch: %d" % (batch_num + 1))
            loss = self.train_batch(train, val, test, idx)
            epoch_loss += loss
            end = time.time()
            self.log.log_scalar('Iteration time', end - start, batch_num + 1 * (self.epochs + 1))
        epoch_loss = epoch_loss / (batch_num + 1)
        self.epochs += 1

        # LOG EPOCH LEVEL INFORMATION
        print('[TRAIN]  Epoch [%d/%d]   Loss: %.4f' % (
            self.epochs, self.max_epochs, epoch_loss))
        info = {'loss': epoch_loss}

        self.log_values(info)
        self.log_hists()

        return epoch_loss

    def train_batch(self, train, val, test, idx):
        self.optimizer.zero_grad()
        network_pred, network_act, _, _, loss_mean_sq_log_diff_level = self.model(train, val,
                                                                                test,
                                                                                idx)

        loss = self.criterion(network_pred, network_act)
        loss.backward()
        nn.utils.clip_grad_value_(self.model.parameters(), self.config['gradient_clipping'])
        self.optimizer.step()
        self.scheduler.step()
        return float(loss)

    def val(self):
        self.model.eval()
        with torch.no_grad():
            acts = []
            preds = []

            hold_out_loss = 0
            for batch_num, (train, val, test, idx) in enumerate(self.dl):
                _, _, (hold_out_pred, network_output_non_train), \
                (hold_out_act, hold_out_act_deseas_norm), _ = self.model(train, val, test, idx)
                hold_out_loss += self.criterion(network_output_non_train.unsqueeze(0).float(),
                                                hold_out_act_deseas_norm.unsqueeze(0).float())
                acts.extend(hold_out_act.view(-1).cpu().detach().numpy())
                preds.extend(hold_out_pred.view(-1).cpu().detach().numpy())

            hold_out_loss = hold_out_loss / (batch_num + 1)

            _hold_out_df = pd.DataFrame({'acts': acts, 'preds': preds})
            #cats = [val for val in self.ohe_headers[info_cat_overall.argmax(axis=1)] for _ in
            #        range(self.config['output_size'])]
            #_hold_out_df['category'] = cats

            overall_hold_out_df = copy.copy(_hold_out_df)
            #overall_hold_out_df['category'] = ['Overall' for _ in cats]

            overall_hold_out_df = pd.concat((_hold_out_df, overall_hold_out_df))
            grouped_results = overall_hold_out_df.groupby(['category']).apply(
                lambda x: np_sMAPE(x.preds, x.acts, x.shape[0]))

            results = grouped_results.to_dict()
            results['hold_out_loss'] = float(hold_out_loss.detach().cpu())

            self.log_values(results)

            file_path = os.path.join('..', 'grouped_results', self.run_id, self.prod_str)
            os.makedirs(file_path, exist_ok=True)

            print(results)
            grouped_path = os.path.join(file_path, 'grouped_results-{}.csv'.format(self.epochs))
            grouped_results.to_csv(grouped_path)
            self.csv_save_path = file_path

        return hold_out_loss.detach().cpu().item()

    def save(self, save_dir='..'):
        print('Loss decreased, saving model!')
        file_path = os.path.join(save_dir, 'models', self.run_id, self.prod_str)
        model_path = os.path.join(file_path, 'model-{}.pyt'.format(self.epochs))
        os.makedirs(file_path, exist_ok=True)
        torch.save({'state_dict': self.model.state_dict()}, model_path)

    def log_values(self, info):

        # SCALAR
        for tag, value in info.items():
            self.log.log_scalar(tag, value, self.epochs + 1)

    def log_hists(self):
        # HISTS
        batch_params = dict()
        for tag, value in self.model.named_parameters():
            if value.grad is not None:
                if "init" in tag:
                    name, _ = tag.split(".")
                    if name not in batch_params.keys() or "%s/grad" % name not in batch_params.keys():
                        batch_params[name] = []
                        batch_params["%s/grad" % name] = []
                    batch_params[name].append(value.data.cpu().numpy())
                    batch_params["%s/grad" % name].append(value.grad.cpu().numpy())
                else:
                    tag = tag.replace('.', '/')
                    self.log.log_histogram(tag, value.data.cpu().numpy(), self.epochs + 1)
                    self.log.log_histogram(tag + '/grad', value.grad.data.cpu().numpy(), self.epochs + 1)
            else:
                print('Not printing %s because it\'s not updating' % tag)

        for tag, v in batch_params.items():
            vals = np.concatenate(np.array(v))
            self.log.log_histogram(tag, vals, self.epochs + 1)

model = ESRNN(num_series=len(dataset), config=config)

tr = ESRNNTrainer(model, dataloader, run_id, config)

tr.train_epochs()

for train, val, test, idx in dataloader:
    print(idx)

# +
#padding problem??
# -

import torch.autograd as autograd

use_cuda = torch.cuda.is_available()

# class DRNN(nn.Module):
#
#     def __init__(self, n_input, n_hidden, n_layers, dilations, dropout=0, cell_type='GRU', batch_first=False):
#
#         super(DRNN, self).__init__()
#
#         self.dilations = dilations
#         self.cell_type = cell_type
#         self.batch_first = batch_first
#
#         layers = []
#         if self.cell_type == "GRU":
#             cell = nn.GRU
#         elif self.cell_type == "RNN":
#             cell = nn.RNN
#         elif self.cell_type == "LSTM":
#             cell = nn.LSTM
#         else:
#             raise NotImplementedError
#
#         for i in range(n_layers):
#             if i == 0:
#                 c = cell(n_input, n_hidden, dropout=dropout)
#             else:
#                 c = cell(n_hidden, n_hidden, dropout=dropout)
#             layers.append(c)
#         self.cells = nn.Sequential(*layers)
#
#     def forward(self, inputs, hidden=None):
#         if self.batch_first:
#             inputs = inputs.transpose(0, 1)
#         outputs = []
#         for i, (cell, dilation) in enumerate(zip(self.cells, self.dilations)): ###### self.dilations
#             if hidden is None:
#                 inputs, _ = self.drnn_layer(cell, inputs, dilation)
#             else:
#                 inputs, hidden[i] = self.drnn_layer(cell, inputs, dilation, hidden[i])
#
#             outputs.append(inputs[-dilation:])
#
#         if self.batch_first:
#             inputs = inputs.transpose(0, 1)
#         return inputs, outputs
#
#     def drnn_layer(self, cell, inputs, rate, hidden=None):
#
#         n_steps = len(inputs)
#         batch_size = inputs[0].size(0)
#         hidden_size = cell.hidden_size
#
#         print("##")
#         print(len(inputs))
#         print(inputs.shape)
#         print('##')
#         inputs, dilated_steps = self._pad_inputs(inputs, n_steps, rate) ## I think this is where the problem is
#         print("###")
#         print(inputs.shape)
#         print('###')
#         dilated_inputs = self._prepare_inputs(inputs, rate) ##
#         print(dilated_inputs.shape)
#         print(torch.eq(inputs, dilated_inputs))
#
#         if hidden is None:
#             dilated_outputs, hidden = self._apply_cell(dilated_inputs, cell, batch_size, rate, hidden_size)
#         else:
#             hidden = self._prepare_inputs(hidden, rate)
#             dilated_outputs, hidden = self._apply_cell(dilated_inputs, cell, batch_size, rate, hidden_size,
#                                                        hidden=hidden)
#
#         splitted_outputs = self._split_outputs(dilated_outputs, rate)
#         outputs = self._unpad_outputs(splitted_outputs, n_steps)
#
#         return outputs, hidden
#
#     def _apply_cell(self, dilated_inputs, cell, batch_size, rate, hidden_size, hidden=None):
#         if hidden is None:
#             if self.cell_type == 'LSTM':
#                 c, m = self.init_hidden(batch_size * rate, hidden_size)
#                 hidden = (c.unsqueeze(0), m.unsqueeze(0))
#             else:
#                 hidden = self.init_hidden(batch_size * rate, hidden_size).unsqueeze(0)
#
#         print(cell)
#         #print()
#         print(hidden[0].shape)
#         print(hidden[1].shape)
#         print(dilated_inputs.shape)
#         dilated_outputs, hidden = cell(dilated_inputs, hidden)
#
#         return dilated_outputs, hidden
#
#     def _unpad_outputs(self, splitted_outputs, n_steps):
#         return splitted_outputs[:n_steps]
#
#     def _split_outputs(self, dilated_outputs, rate):
#         batchsize = dilated_outputs.size(1) // rate
#
#         blocks = [dilated_outputs[:, i * batchsize: (i + 1) * batchsize, :] for i in range(rate)]
#
#         interleaved = torch.stack((blocks)).transpose(1, 0).contiguous()
#         interleaved = interleaved.view(dilated_outputs.size(0) * rate,
#                                        batchsize,
#                                        dilated_outputs.size(2))
#         return interleaved
#
#     def _pad_inputs(self, inputs, n_steps, rate):
#         iseven = (n_steps % rate) == 0 ###########
#
#         if not iseven:
#             dilated_steps = n_steps // rate + 1
#
#             zeros_ = torch.zeros(dilated_steps * rate - inputs.size(0),
#                                  inputs.size(1),
#                                  inputs.size(2))
#             if use_cuda:
#                 zeros_ = zeros_.cuda()
#
#             inputs = torch.cat((inputs, autograd.Variable(zeros_)))
#         else:
#             dilated_steps = n_steps // rate
#
#         return inputs, dilated_steps
#
#     def _prepare_inputs(self, inputs, rate):
#         dilated_inputs = torch.cat([inputs[j::rate, :, :] for j in range(rate)], 1)
#         return dilated_inputs
#
#     def init_hidden(self, batch_size, hidden_dim):
#         hidden = autograd.Variable(torch.zeros(batch_size, hidden_dim))
#         if use_cuda:
#             hidden = hidden.cuda()
#         if self.cell_type == "LSTM":
#             memory = autograd.Variable(torch.zeros(batch_size, hidden_dim))
#             if use_cuda:
#                 memory = memory.cuda()
#             return hidden, memory
#         else:
#             return hidden

# num_series = 1
#
# init_lev_sms = []
# init_seas_sms = []
# init_seasonalities = []
#
# for i in range(num_series):
#     init_lev_sms.append(nn.Parameter(torch.Tensor([0.5]), requires_grad=True))
#     init_seas_sms.append(nn.Parameter(torch.Tensor([0.5]), requires_grad=True))
#     init_seasonalities.append(nn.Parameter((torch.ones(config['seasonality']) * 0.5), requires_grad=True))
#
# init_seasonalities
#
# nn.ParameterList(init_lev_sms)
#
# init_lev_sms = nn.ParameterList(init_lev_sms)
# init_seas_sms = nn.ParameterList(init_seas_sms)
# init_seasonalities = nn.ParameterList(init_seasonalities)
#
# init_seasonalities
#
# idxs = [torch.tensor(0)]
#
# [init_lev_sms[idx] for idx in idxs]
#
# nn.Sigmoid()(torch.stack([init_lev_sms[idx] for idx in idxs]).squeeze(1))
#
# logistic = nn.Sigmoid()
#
# lev_sms = logistic(torch.stack([init_lev_sms[idx] for idx in idxs]).squeeze(1))
# seas_sms = logistic(torch.stack([init_seas_sms[idx] for idx in idxs]).squeeze(1))
# init_seasonalities = torch.stack([init_seasonalities[idx] for idx in idxs])
#
# seasonalities = []
# # PRIME SEASONALITY
# for i in range(config['seasonality']):
#     seasonalities.append(torch.exp(init_seasonalities[:, i]))
# seasonalities.append(torch.exp(init_seasonalities[:, 0]))
#
# seasonalities
#
# levs = []
# log_diff_of_levels = []
#
# levs.append(train[:, 0] / seasonalities[0])
# for i in range(1, train.shape[1]):
#     # CALCULATE LEVEL FOR CURRENT TIMESTEP TO NORMALIZE RNN
#     new_lev = lev_sms * (train[:, i] / seasonalities[i]) + (1 - lev_sms) * levs[i - 1]
#     levs.append(new_lev)
#
#     # STORE DIFFERENCE TO PENALIZE LATER
#     log_diff_of_levels.append(torch.log(new_lev / levs[i - 1]))
#
#     # CALCULATE SEASONALITY TO DESEASONALIZE THE DATA FOR RNN
#     seasonalities.append(seas_sms * (train[:, i] / new_lev) + (1 - seas_sms) * seasonalities[i])
#
# seasonalities_stacked = torch.stack(seasonalities).transpose(1, 0)
# levs_stacked = torch.stack(levs).transpose(1, 0)
#
# loss_mean_sq_log_diff_level = 0
# if config['level_variability_penalty'] > 0:
#     sq_log_diff = torch.stack(
#         [(log_diff_of_levels[i] - log_diff_of_levels[i - 1]) ** 2 for i in range(1, len(log_diff_of_levels))])
#     loss_mean_sq_log_diff_level = torch.mean(sq_log_diff)
#
# if config['output_size'] > config['seasonality']:
#     start_seasonality_ext = seasonalities_stacked.shape[1] - config['seasonality']
#     end_seasonality_ext = start_seasonality_ext + config['output_size'] - config['seasonality']
#     seasonalities_stacked = torch.cat((seasonalities_stacked, seasonalities_stacked[:, start_seasonality_ext:end_seasonality_ext]),
#                                       dim=1)
#
#
#


text:
# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/shapley_values.ipynb (unless otherwise specified).

__all__ = ['ShapleyValues']

# Cell
# Author: Simon Grah <simon.grah@thalesgroup.com>
#         Vincent Thouvenot <vincent.thouvenot@thalesgroup.com>

# MIT License

# Copyright (c) 2020 Thales Six GTS France

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# Cell
import numpy as np
import pandas as pd
from itertools import combinations
from math import factorial
from tqdm import tqdm

# Cell
def ShapleyValues(x, fc, ref):
    """
    Calculate the exact Shapley Values for an individual x
    in a game based on a reference r and the reward function fc.
    """

    # Get general information
    feature_names = list(x.index)
    d = len(feature_names) # dimension
    set_features = set(feature_names)

    # Store Shapley Values in a pandas Series
    Φ = pd.Series(np.zeros(d), index=feature_names)

    # Individual reference or dataset of references
    def output_single_ref(coalition, feature_names):
        z = np.array([x[col] if col in coalition else ref.loc[col] for col in feature_names])
        return fc(z)

    def output_several_ref(coalition, feature_names):
        rewards = []
        idxs = np.random.choice(ref.index, size=len(ref), replace=False)
        for idx in idxs:
            z = np.array([x[col] if col in coalition else ref.loc[idx, col] for col in feature_names])
            rewards.append(fc(z))
        return np.mean(rewards)

    if isinstance(ref, pd.core.series.Series):
        individual_ref = True
        output = output_single_ref
    elif isinstance(ref, pd.core.frame.DataFrame):
        if ref.shape[0] == 1:
            ref = ref.iloc[0]
            individual_ref = True
            output = output_single_ref
        else:
            individual_ref = False
            output = output_several_ref

    # Start computation (number of coalitions: 2**d - 1)
    for cardinal_S in tqdm(range(0, d)):
        # weight
        ω = factorial(cardinal_S) * (factorial(d - cardinal_S - 1))
        ω /= factorial(d)
        # iter over all combinations of size cardinal_S
        for S in combinations(feature_names, cardinal_S):
            S = list(S)
            f_S = output(S, feature_names)
            # Consider only features outside of S
            features_out_S = set_features - set(S)
            for j in features_out_S:
                S_union_j = S + [j]
                f_S_union_j = output(S_union_j, feature_names)
                # Update Shapley value of attribute i
                Φ[j] += ω * (f_S_union_j - f_S)

    return Φ

max_stars_repo_name: alheliou/shapkit

max_issues_repo_name: alheliou/shapkit

max_forks_repo_name: alheliou/shapkit


text:
import numpy as np
from pmesh.pm     import ParticleMesh
from nbodykit.lab import BigFileCatalog, MultipleSpeciesCatalog,\
                         BigFileMesh, FFTPower
from nbodykit     import setup_logging
from mpi4py       import MPI

import HImodels
# enable logging, we have some clue what's going on.
setup_logging('info')

#
#Global, fixed things
scratchyf = '/global/cscratch1/sd/yfeng1/m3127/'
scratchcm = '/global/cscratch1/sd/chmodi/m3127/H1mass/'
project  = '/project/projectdirs/m3127/H1mass/'
cosmodef = {'omegam':0.309167, 'h':0.677, 'omegab':0.048}
alist    = [0.1429,0.1538,0.1667,0.1818,0.2000,0.2222,0.2500,0.2857,0.3333]

#Parameters, box size, number of mesh cells, simulation, ...
bs, nc, ncsim, sim, prefix = 256, 512, 2560, 'highres/%d-9100-fixed'%2560, 'highres'
#bs,nc,ncsim, sim, prefic = 1024, 1024, 10240, 'highres/%d-9100-fixed'%ncsim, 'highres'

# It's useful to have my rank for printing...
pm   = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
rank = pm.comm.rank
comm = pm.comm

#Which model to use
HImodel = HImodels.ModelA
ofolder = '../data/outputs/'

def calc_bias(aa,h1mesh,suff):
    '''Compute the bias(es) for the HI'''
    if rank==0:
        print("Processing a={:.4f}...".format(aa))
        print('Reading DM mesh...')
    if ncsim == 10240:
        dm    = BigFileMesh(scratchyf+sim+'/fastpm_%0.4f/'%aa+\
                            '/1-mesh/N%04d'%nc,'').paint()
    else:
        dm    = BigFileMesh(project+sim+'/fastpm_%0.4f/'%aa+\
                        '/dmesh_N%04d/1/'%nc,'').paint()
    dm   /= dm.cmean()
    if rank==0: print('Computing DM P(k)...')
    pkmm  = FFTPower(dm,mode='1d').power
    k,pkmm= pkmm['k'],pkmm['power']  # Ignore shotnoise.
    if rank==0: print('Done.')
    #

    pkh1h1 = FFTPower(h1mesh,mode='1d').power
    pkh1h1 = pkh1h1['power']-pkh1h1.attrs['shotnoise']
    pkh1mm = FFTPower(h1mesh,second=dm,mode='1d').power['power']
    if rank==0: print('Done.')
    # Compute the biases.
    b1x = np.abs(pkh1mm/(pkmm+1e-10))
    b1a = np.abs(pkh1h1/(pkmm+1e-10))**0.5
    if rank==0: print("Finishing processing a={:.4f}.".format(aa))
    return(k,b1x,b1a,np.abs(pkmm))
    #

if __name__=="__main__":
    #satsuff='-m1_5p0min-alpha_0p8-16node'
    suff='-m1_00p3mh-alpha-0p8-subvol'
    outfolder = ofolder + suff[1:] + "/modelA/"
    try:  os.makedirs(outfolder)
    except : pass

    if rank==0:
        print('Starting')
    for aa in alist:

        if rank == 0: print('\n ############## Redshift = %0.2f ############## \n'%(1/aa-1))
        halocat = BigFileCatalog(scratchyf + sim+ '/fastpm_%0.4f//'%aa, dataset='LL-0.200')
        mp = halocat.attrs['MassTable'][1]*1e10##
        halocat['Mass'] = halocat['Length'].compute() * mp
        cencat = BigFileCatalog(scratchcm + sim+'/fastpm_%0.4f/cencat'%aa+suff)
        satcat = BigFileCatalog(scratchcm + sim+'/fastpm_%0.4f/satcat'%aa+suff)
        rsdfac = read_conversions(scratchyf + sim+'/fastpm_%0.4f/'%aa)
        #

        HImodelz = HImodel(aa)
        los = [0,0,1]
        halocat['HImass'], cencat['HImass'], satcat['HImass'] = HImodelz.assignHI(halocat, cencat, satcat)
        halocat['RSDpos'], cencat['RSDpos'], satcat['RSDpos'] = HImodelz.assignrsd(rsdfac, halocat, cencat, satcat, los=los)

        h1mesh = HImodelz.createmesh(bs, nc, halocat, cencat, satcat, mode='galaxies', position='RSDpos', weight='HImass')

        kk,b1x,b1a,pkmm = calc_bias(aa,h1mesh,suff)

        #
        if rank==0:
            fout = open(outfolder + "HI_bias_{:6.4f}.txt".format(aa),"w")
            fout.write("# Mcut={:12.4e}Msun/h.\n".format(mcut))
            fout.write("# {:>8s} {:>10s} {:>10s} {:>15s}\n".\
                       format("k","b1_x","b1_a","Pkmm"))
            for i in range(1,kk.size):
                fout.write("{:10.5f} {:10.5f} {:10.5f} {:15.5e}\n".\
                           format(kk[i],b1x[i],b1a[i],pkmm[i]))
            fout.close()
    #


text:
#!/usr/bin/env python3 -u
# -*- coding: utf-8 -*-

__author__ = ["Markus Löning"]
__all__ = [
    "TEST_YS",
    "TEST_SPS",
    "TEST_ALPHAS",
    "TEST_FHS",
    "TEST_STEP_LENGTHS_INT",
    "TEST_STEP_LENGTHS",
    "TEST_INS_FHS",
    "TEST_OOS_FHS",
    "TEST_WINDOW_LENGTHS_INT",
    "TEST_WINDOW_LENGTHS",
    "TEST_INITIAL_WINDOW_INT",
    "TEST_INITIAL_WINDOW",
    "VALID_INDEX_FH_COMBINATIONS",
    "INDEX_TYPE_LOOKUP",
    "TEST_RANDOM_SEEDS",
    "TEST_N_ITERS",
]

import numpy as np
import pandas as pd

from sktime.utils._testing.series import _make_series

# We here define the parameter values for unit testing.
TEST_CUTOFFS_INT = [np.array([21, 22]), np.array([3, 7, 10])]
# The following timestamps correspond
# to the above integers for `_make_series(all_positive=True)`
TEST_CUTOFFS_TIMESTAMP = [
    pd.to_datetime(["2000-01-22", "2000-01-23"]),
    pd.to_datetime(["2000-01-04", "2000-01-08", "2000-01-11"]),
]
TEST_CUTOFFS = [*TEST_CUTOFFS_INT, *TEST_CUTOFFS_TIMESTAMP]

TEST_WINDOW_LENGTHS_INT = [1, 5]
TEST_WINDOW_LENGTHS_TIMEDELTA = [pd.Timedelta(1, unit="D"), pd.Timedelta(5, unit="D")]
TEST_WINDOW_LENGTHS_DATEOFFSET = [pd.offsets.Day(1), pd.offsets.Day(5)]
TEST_WINDOW_LENGTHS = [
    *TEST_WINDOW_LENGTHS_INT,
    *TEST_WINDOW_LENGTHS_TIMEDELTA,
    *TEST_WINDOW_LENGTHS_DATEOFFSET,
]

TEST_INITIAL_WINDOW_INT = [7, 10]
TEST_INITIAL_WINDOW_TIMEDELTA = [pd.Timedelta(7, unit="D"), pd.Timedelta(10, unit="D")]
TEST_INITIAL_WINDOW_DATEOFFSET = [pd.offsets.Day(7), pd.offsets.Day(10)]
TEST_INITIAL_WINDOW = [
    *TEST_INITIAL_WINDOW_INT,
    *TEST_INITIAL_WINDOW_TIMEDELTA,
    *TEST_INITIAL_WINDOW_DATEOFFSET,
]

TEST_STEP_LENGTHS_INT = [1, 5]
TEST_STEP_LENGTHS_TIMEDELTA = [pd.Timedelta(1, unit="D"), pd.Timedelta(5, unit="D")]
TEST_STEP_LENGTHS_DATEOFFSET = [pd.offsets.Day(1), pd.offsets.Day(5)]
TEST_STEP_LENGTHS = [
    *TEST_STEP_LENGTHS_INT,
    *TEST_STEP_LENGTHS_TIMEDELTA,
    *TEST_STEP_LENGTHS_DATEOFFSET,
]

TEST_OOS_FHS = [1, np.array([2, 5], dtype="int64")]  # out-of-sample
TEST_INS_FHS = [
    -3,  # single in-sample
    np.array([-2, -5], dtype="int64"),  # multiple in-sample
    0,  # last training point
    np.array([-3, 2], dtype="int64"),  # mixed in-sample and out-of-sample
]
TEST_FHS = [*TEST_OOS_FHS, *TEST_INS_FHS]

TEST_OOS_FHS_TIMEDELTA = [
    [pd.Timedelta(1, unit="D")],
    [pd.Timedelta(2, unit="D"), pd.Timedelta(5, unit="D")],
]  # out-of-sample
TEST_INS_FHS_TIMEDELTA = [
    pd.Timedelta(-3, unit="D"),  # single in-sample
    [pd.Timedelta(-2, unit="D"), pd.Timedelta(-5, unit="D")],  # multiple in-sample
    pd.Timedelta(0, unit="D"),  # last training point
    [
        pd.Timedelta(-3, unit="D"),
        pd.Timedelta(2, unit="D"),
    ],  # mixed in-sample and out-of-sample
]
TEST_FHS_TIMEDELTA = [*TEST_OOS_FHS_TIMEDELTA, *TEST_INS_FHS_TIMEDELTA]

TEST_SPS = [3, 12]
TEST_ALPHAS = [0.05, 0.1]
TEST_YS = [_make_series(all_positive=True)]
TEST_RANDOM_SEEDS = [1, 42]
TEST_N_ITERS = [1, 4]

# We currently support the following combinations of index and forecasting horizon types
VALID_INDEX_FH_COMBINATIONS = [
    # index type, fh type, is_relative
    ("int", "int", True),
    ("int", "int", False),
    ("range", "int", True),
    ("range", "int", False),
    ("period", "int", True),
    ("period", "period", False),
    ("datetime", "int", True),
    ("datetime", "datetime", False),
    ("datetime", "timedelta", True),
]

INDEX_TYPE_LOOKUP = {
    "int": pd.Index,
    "range": pd.RangeIndex,
    "datetime": pd.DatetimeIndex,
    "period": pd.PeriodIndex,
    "timedelta": pd.TimedeltaIndex,
}

max_stars_repo_name: biologioholic/sktime

max_issues_repo_name: biologioholic/sktime

max_forks_repo_name: biologioholic/sktime


text:
import argparse
import numpy as np
import cv2
import os
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader

from model import FFDNet
import utils

def read_image(image_path, is_gray):
    """
    :return: Normalized Image (C * W * H)
    """
    if is_gray:
        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        image = np.expand_dims(image.T, 0) # 1 * W * H
    else:
        image = cv2.imread(image_path)
        image = (cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).transpose(2, 1, 0) # 3 * W * H

    return utils.normalize(image)

def load_images(is_train, is_gray, base_path):
    """
    :param base_path: ./train_data/
    :return: List[Patches] (C * W * H)
    """
    if is_gray:
        train_dir = 'gray/train/'
        val_dir = 'gray/val/'
    else:
        train_dir = 'rgb/train/'
        val_dir = 'rgb/val/'

    image_dir = base_path.replace('\'', '').replace('"', '') + (train_dir if is_train else val_dir)
    print('> Loading images in ' + image_dir)
    images = []
    for fn in next(os.walk(image_dir))[2]:
        image = read_image(image_dir + fn, is_gray)
        images.append(image)
    return images

def images_to_patches(images, patch_size):
    """
    :param images: List[Image (C * W * H)]
    :param patch_size: int
    :return: (n * C * W * H)
    """
    patches_list = []
    for image in images:
        patches = utils.image_to_patches(image, patch_size=patch_size)
        if len(patches) != 0:
            patches_list.append(patches)
    del images
    return np.vstack(patches_list)

def train(args):
    print('> Loading dataset...')
    # Images
    train_dataset = load_images(is_train=True, is_gray=args.is_gray, base_path=args.train_path)
    val_dataset = load_images(is_train=False, is_gray=args.is_gray, base_path=args.train_path)
    print(f'\tTrain image datasets: {len(train_dataset)}')
    print(f'\tVal image datasets: {len(val_dataset)}')

    # Patches
    train_dataset = images_to_patches(train_dataset, patch_size=args.patch_size)
    val_dataset = images_to_patches(val_dataset, patch_size=args.patch_size)
    print(f'\tTrain patch datasets: {train_dataset.shape}')
    print(f'\tVal patch datasets: {val_dataset.shape}')

    # DataLoader
    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=6)
    val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=6)
    print(f'\tTrain batch number: {len(train_dataloader)}')
    print(f'\tVal batch number: {len(val_dataloader)}')

    # Noise list
    train_noises = args.train_noise_interval # [0, 75, 15]
    val_noises = args.val_noise_interval # [0, 60, 30]
    train_noises = list(range(train_noises[0], train_noises[1], train_noises[2]))
    val_noises = list(range(val_noises[0], val_noises[1], val_noises[2]))
    print(f'\tTrain noise internal: {train_noises}')
    print(f'\tVal noise internal: {val_noises}')
    print('\n')

    # Model & Optim
    model = FFDNet(is_gray=args.is_gray)
    model.apply(utils.weights_init_kaiming)
    if args.cuda:
        model = model.cuda()
    loss_fn = nn.MSELoss(reduction='sum')
    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)

    print('> Start training...')
    for epoch_idx in range(args.epoches):
        # Train
        loss_idx = 0
        train_losses = 0
        model.train()

        start_time = time.time()
        for batch_idx, batch_data in enumerate(train_dataloader):
            # According to internal, add noise
            for int_noise_sigma in train_noises:
                noise_sigma = int_noise_sigma / 255
                new_images = utils.add_batch_noise(batch_data, noise_sigma)
                noise_sigma = torch.FloatTensor(np.array([noise_sigma for idx in range(new_images.shape[0])]))
                new_images = Variable(new_images)
                noise_sigma = Variable(noise_sigma)
                if args.cuda:
                    new_images = new_images.cuda()
                    noise_sigma = noise_sigma.cuda()

                # Predict
                images_pred = model(new_images, noise_sigma)
                train_loss = loss_fn(images_pred, batch_data.to(images_pred.device))
                train_losses += train_loss
                loss_idx += 1

                optimizer.zero_grad()
                train_loss.backward()
                optimizer.step()

                # Log Progress
                stop_time = time.time()
                all_num = len(train_dataloader) * len(train_noises)
                done_num = batch_idx * len(train_noises) + train_noises.index(int_noise_sigma) + 1
                rest_time = int((stop_time - start_time) / done_num * (all_num - done_num))
                percent = int(done_num / all_num * 100)
                print(f'\rEpoch: {epoch_idx + 1} / {args.epoches}, ' +
                      f'Batch: {batch_idx + 1} / {len(train_dataloader)}, ' +
                      f'Noise_Sigma: {int_noise_sigma} / {train_noises[-1]}, ' +
                      f'Train_Loss: {train_loss}, ' +
                      f'=> {rest_time}s, {percent}%', end='')

        train_losses /= loss_idx
        print(f', Avg_Train_Loss: {train_losses}, All: {int(stop_time - start_time)}s')

        # Evaluate
        loss_idx = 0
        val_losses = 0
        if (epoch_idx + 1) % args.val_epoch != 0:
            continue
        model.eval()

        start_time = time.time()
        for batch_idx, batch_data in enumerate(val_dataloader):
            # According to internal, add noise
            for int_noise_sigma in val_noises:
                noise_sigma = int_noise_sigma / 255
                new_images = utils.add_batch_noise(batch_data, noise_sigma)
                noise_sigma = torch.FloatTensor(np.array([noise_sigma for idx in range(new_images.shape[0])]))
                new_images = Variable(new_images)
                noise_sigma = Variable(noise_sigma)
                if args.cuda:
                    new_images = new_images.cuda()
                    noise_sigma = noise_sigma.cuda()

                # Predict
                images_pred = model(new_images, noise_sigma)
                val_loss = loss_fn(images_pred, batch_data.to(images_pred.device))
                val_losses += val_loss
                loss_idx += 1

                # Log Progress
                stop_time = time.time()
                all_num = len(val_dataloader) * len(val_noises)
                done_num = batch_idx * len(val_noises) + val_noises.index(int_noise_sigma) + 1
                rest_time = int((stop_time - start_time) / done_num * (all_num - done_num))
                percent = int(done_num / all_num * 100)
                print(f'\rEpoch: {epoch_idx + 1} / {args.epoches}, ' +
                      f'Batch: {batch_idx + 1} / {len(val_dataloader)}, ' +
                      f'Noise_Sigma: {int_noise_sigma} / {val_noises[-1]}, ' +
                      f'Val_Loss: {val_loss}, ' +
                      f'=> {rest_time}s, {percent}%', end='')

        val_losses /= loss_idx
        print(f', Avg_Val_Loss: {val_losses}, All: {int(stop_time - start_time)}s')

        # Save Checkpoint
        if (epoch_idx + 1) % args.save_checkpoints == 0:
            model_path = args.model_path + ('net_gray_checkpoint.pth' if args.is_gray else 'net_rgb_checkpoint.pth')
            torch.save(model.state_dict(), model_path)
            print(f'| Saved Checkpoint at Epoch {epoch_idx + 1} to {model_path}')

    # Final Save Model Dict
    model.eval()
    model_path = args.model_path + ('net_gray.pth' if args.is_gray else 'net_rgb.pth')
    torch.save(model.state_dict(), model_path)
    print(f'Saved State Dict in {model_path}')
    print('\n')

def test(args):
    # Image
    image = cv2.imread(args.test_path)
    if image is None:
        raise Exception(f'File {args.test_path} not found or error')
    is_gray = utils.is_image_gray(image)
    image = read_image(args.test_path, is_gray)
    print("{} image shape: {}".format("Gray" if is_gray else "RGB", image.shape))

    # Expand odd shape to even
    expend_W = False
    expend_H = False
    if image.shape[1] % 2 != 0:
        expend_W = True
        image = np.concatenate((image, image[:, -1, :][:, np.newaxis, :]), axis=1)
    if image.shape[2] % 2 != 0:
        expend_H = True
        image = np.concatenate((image, image[:, :, -1][:, :, np.newaxis]), axis=2)

    # Noise
    image = torch.FloatTensor([image]) # 1 * C(1 / 3) * W * H
    if args.add_noise:
        image = utils.add_batch_noise(image, args.noise_sigma)
    noise_sigma = torch.FloatTensor([args.noise_sigma])

    # Model & GPU
    model = FFDNet(is_gray=is_gray)
    if args.cuda:
        image = image.cuda()
        noise_sigma = noise_sigma.cuda()
        model = model.cuda()

    # Dict
    model_path = args.model_path + ('net_gray.pth' if is_gray else 'net_rgb.pth')
    print(f"> Loading model param in {model_path}...")
    state_dict = torch.load(model_path)
    model.load_state_dict(state_dict)
    model.eval()
    print('\n')

    # Test
    with torch.no_grad():
        start_time = time.time()
        image_pred = model(image, noise_sigma)
        stop_time = time.time()
        print("Test time: {0:.4f}s".format(stop_time - start_time))

    # PSNR
    psnr = utils.batch_psnr(img=image_pred, imclean=image, data_range=1)
    print("PSNR denoised {0:.2f}dB".format(psnr))

    # UnExpand odd
    if expend_W:
        image_pred = image_pred[:, :, :-1, :]
    if expend_H:
        image_pred = image_pred[:, :, :, :-1]

    # Save
    cv2.imwrite("ffdnet.png", utils.variable_to_cv2_image(image_pred))
    if args.add_noise:
        cv2.imwrite("noisy.png", utils.variable_to_cv2_image(image))

def main():
    parser = argparse.ArgumentParser()

    # Train
    parser.add_argument("--train_path", type=str, default='./train_data/',                  help='Train dataset dir.')
    parser.add_argument("--is_gray", action='store_true',                                   help='Train gray/rgb model.')
    parser.add_argument("--patch_size", type=int, default=32,                               help='Uniform size of training images patches.')
    parser.add_argument("--train_noise_interval", nargs=3, type=int, default=[0, 75, 15],   help='Train dataset noise sigma set interval.')
    parser.add_argument("--val_noise_interval", nargs=3, type=int, default=[0, 60, 30],     help='Validation dataset noise sigma set interval.')
    parser.add_argument("--batch_size", type=int, default=256,                              help='Batch size for training.')
    parser.add_argument("--epoches", type=int, default=80,                                  help='Total number of training epoches.')
    parser.add_argument("--val_epoch", type=int, default=5,                                 help='Total number of validation epoches.')
    parser.add_argument("--learning_rate", type=float, default=1e-3,                        help='The initial learning rate for Adam.')
    parser.add_argument("--save_checkpoints", type=int, default=5,                          help='Save checkpoint every epoch.')

    # Test
    parser.add_argument("--test_path", type=str, default='./test_data/color.png',           help='Test image path.')
    parser.add_argument("--noise_sigma", type=float, default=25,                            help='Input uniform noise sigma for test.')
    parser.add_argument('--add_noise', action='store_true',                                 help='Add noise_sigma to input or not.')

    # Global
    parser.add_argument("--model_path", type=str, default='./models/',                      help='Model loading and saving path.')
    parser.add_argument("--use_gpu", action='store_true',                                   help='Train and test using GPU.')
    parser.add_argument("--is_train", action='store_true',                                  help='Do train.')
    parser.add_argument("--is_test", action='store_true',                                   help='Do test.')

    args = parser.parse_args()
    assert (args.is_train or args.is_test), 'is_train 和 is_test 至少有一个为 True'

    args.cuda = args.use_gpu and torch.cuda.is_available()
    print("> Parameters: ")
    for k, v in zip(args.__dict__.keys(), args.__dict__.values()):
        print(f'\t{k}: {v}')
    print('\n')

    # Normalize noise level
    args.noise_sigma /= 255
    args.train_noise_interval[1] += 1
    args.val_noise_interval[1] += 1

    if args.is_train:
        train(args)

    if args.is_test:
        test(args)

if __name__ == "__main__":
    main()


text:
""" Unit tests for the problem interface."""

import unittest
import warnings
from six import assertRaisesRegex

import numpy as np

from openmdao.api import Problem, Group, IndepVarComp, PETScVector, NonlinearBlockGS, \
     ScipyOptimizer
from openmdao.devtools.testutil import assert_rel_error
from openmdao.test_suite.components.paraboloid import Paraboloid

from openmdao.test_suite.components.sellar import SellarDerivatives, SellarDerivativesConnected

class TestProblem(unittest.TestCase):

    def test_set_2d_array(self):

        prob = Problem(model=Group())
        model = prob.model
        model.add_subsystem(name='indeps',
                            subsys=IndepVarComp(name='X_c', shape=(3, 1)))
        prob.setup()

        new_val = -5*np.ones((3, 1))
        prob['indeps.X_c'] = new_val
        assert_rel_error(self, prob['indeps.X_c'], new_val, 1e-10)

        new_val = 2.5*np.ones(3)
        prob['indeps.X_c'][:, 0] = new_val
        assert_rel_error(self, prob['indeps.X_c'], new_val.reshape((3,1)), 1e-10)
        assert_rel_error(self, prob['indeps.X_c'][:, 0], new_val, 1e-10)

    def test_set_checks_shape(self):

        model = Group()

        indep = model.add_subsystem('indep', IndepVarComp())
        indep.add_output('num')
        indep.add_output('arr', shape=(10, 1))

        prob = Problem(model)
        prob.setup()

        msg = "Incompatible shape for '.*': Expected (.*) but got (.*)"

        # check valid scalar value
        new_val = -10.
        prob['indep.num'] = new_val
        assert_rel_error(self, prob['indep.num'], new_val, 1e-10)

        # check bad scalar value
        bad_val = -10*np.ones((10))
        with assertRaisesRegex(self, ValueError, msg):
            prob['indep.num'] = bad_val

        # check assign scalar to array
        arr_val = new_val*np.ones((10, 1))
        prob['indep.arr'] = new_val
        assert_rel_error(self, prob['indep.arr'], arr_val, 1e-10)

        # check valid array value
        new_val = -10*np.ones((10, 1))
        prob['indep.arr'] = new_val
        assert_rel_error(self, prob['indep.arr'], new_val, 1e-10)

        # check bad array value
        bad_val = -10*np.ones((10))
        with assertRaisesRegex(self, ValueError, msg):
            prob['indep.arr'] = bad_val

        # check valid list value
        new_val = new_val.tolist()
        prob['indep.arr'] = new_val
        assert_rel_error(self, prob['indep.arr'], new_val, 1e-10)

        # check bad list value
        bad_val = bad_val.tolist()
        with assertRaisesRegex(self, ValueError, msg):
            prob['indep.arr'] = bad_val

    def test_compute_total_derivs_basic(self):
        # Basic test for the method using default solvers on simple model.

        prob = Problem()
        model = prob.model = Group()
        model.add_subsystem('p1', IndepVarComp('x', 0.0), promotes=['x'])
        model.add_subsystem('p2', IndepVarComp('y', 0.0), promotes=['y'])
        model.add_subsystem('comp', Paraboloid(), promotes=['x', 'y', 'f_xy'])

        prob.setup(check=False, mode='fwd')
        prob.set_solver_print(level=0)
        prob.run_model()

        of = ['f_xy']
        wrt = ['x', 'y']
        derivs = prob.compute_total_derivs(of=of, wrt=wrt)

        assert_rel_error(self, derivs['f_xy', 'x'], [[-6.0]], 1e-6)
        assert_rel_error(self, derivs['f_xy', 'y'], [[8.0]], 1e-6)

        prob.setup(check=False, mode='rev')
        prob.run_model()

        of = ['f_xy']
        wrt = ['x', 'y']
        derivs = prob.compute_total_derivs(of=of, wrt=wrt)

        assert_rel_error(self, derivs['f_xy', 'x'], [[-6.0]], 1e-6)
        assert_rel_error(self, derivs['f_xy', 'y'], [[8.0]], 1e-6)

    def test_compute_total_derivs_basic_return_dict(self):
        # Make sure 'dict' return_format works.

        prob = Problem()
        model = prob.model = Group()
        model.add_subsystem('p1', IndepVarComp('x', 0.0), promotes=['x'])
        model.add_subsystem('p2', IndepVarComp('y', 0.0), promotes=['y'])
        model.add_subsystem('comp', Paraboloid(), promotes=['x', 'y', 'f_xy'])

        prob.setup(check=False, mode='fwd')
        prob.set_solver_print(level=0)
        prob.run_model()

        of = ['f_xy']
        wrt = ['x', 'y']
        derivs = prob.compute_total_derivs(of=of, wrt=wrt, return_format='dict')

        assert_rel_error(self, derivs['f_xy']['x'], [[-6.0]], 1e-6)
        assert_rel_error(self, derivs['f_xy']['y'], [[8.0]], 1e-6)

        prob.setup(check=False, mode='rev')
        prob.run_model()

        of = ['f_xy']
        wrt = ['x', 'y']
        derivs = prob.compute_total_derivs(of=of, wrt=wrt, return_format='dict')

        assert_rel_error(self, derivs['f_xy']['x'], [[-6.0]], 1e-6)
        assert_rel_error(self, derivs['f_xy']['y'], [[8.0]], 1e-6)

    def test_feature_set_indeps(self):
        prob = Problem()

        model = prob.model = Group()
        model.add_subsystem('p1', IndepVarComp('x', 0.0), promotes=['x'])
        model.add_subsystem('p2', IndepVarComp('y', 0.0), promotes=['y'])
        model.add_subsystem('comp', Paraboloid(), promotes=['x', 'y', 'f_xy'])

        prob.setup()

        prob['x'] = 2.
        prob['y'] = 10.
        prob.run_model()
        assert_rel_error(self, prob['f_xy'], 214.0, 1e-6)

    def test_feature_numpyvec_setup(self):

        prob = Problem()
        model = prob.model = Group()
        model.add_subsystem('p1', IndepVarComp('x', 0.0), promotes=['x'])
        model.add_subsystem('p2', IndepVarComp('y', 0.0), promotes=['y'])
        model.add_subsystem('comp', Paraboloid(), promotes=['x', 'y', 'f_xy'])

        prob.setup()

        prob['x'] = 2.
        prob['y'] = 10.
        prob.run_model()
        assert_rel_error(self, prob['f_xy'], 214.0, 1e-6)

        prob['x'] = 0.
        prob['y'] = 0.
        prob.run_model()
        assert_rel_error(self, prob['f_xy'], 22.0, 1e-6)

        # skip the setup error checking
        prob.setup(check=False)
        prob['x'] = 4
        prob['y'] = 8.

        prob.run_model()
        assert_rel_error(self, prob['f_xy'], 174.0, 1e-6)

    @unittest.skipUnless(PETScVector, "PETSc is required.")
    def test_feature_petsc_setup(self):

        prob = Problem()
        model = prob.model = Group()
        model.add_subsystem('p1', IndepVarComp('x', 0.0), promotes=['x'])
        model.add_subsystem('p2', IndepVarComp('y', 0.0), promotes=['y'])
        model.add_subsystem('comp', Paraboloid(), promotes=['x', 'y', 'f_xy'])

        # use PETScVector when using any PETSc linear solvers or running under MPI
        prob.setup(vector_class=PETScVector)
        prob['x'] = 2.
        prob['y'] = 10.

        prob.run_model()
        assert_rel_error(self, prob['f_xy'], 214.0, 1e-6)

    def test_feature_check_total_derivatives_manual(self):
        raise unittest.SkipTest("check_total_derivatives not implemented yet")

        prob = Problem()
        prob.model = SellarDerivatives()
        prob.model.nonlinear_solver = NonlinearBlockGS()

        prob.setup()
        prob.run_model()

        # manually specify which derivatives to check
        # TODO: need a decorator to capture this output and put it into the doc,
        #       or maybe just a new kind of assert?
        prob.check_total_derivatives(of=['obj', 'con1'], wrt=['x', 'z'])
        # TODO: Need to devlop the group FD/CS api, so user can control how this
        #       happens by chaninging settings on the root node

    def test_feature_check_total_derivatives_from_driver(self):
        raise unittest.SkipTest("check_total_derivatives not implemented yet")

        prob = Problem()
        prob.model = SellarDerivatives()
        prob.model.nonlinear_solver = NonlinearBlockGS()

        prob.setup()

        prob.model.options['method'] = 'slsqp'
        prob.model.add_design_var('x', lower=-100, upper=100)
        prob.model.add_design_var('z', lower=-100, upper=100)
        prob.model.add_objective('obj')
        prob.model.add_design_var('con1')
        prob.model.add_design_var('con2')
        # re-do setup since we changed the driver and problem inputs/outputs
        prob.setup()

        # We don't call run_driver() here because we don't
        # actually want the optimizer to run
        prob.run_model()

        # check derivatives of all obj+constraints w.r.t all design variables
        prob.check_total_derivatives()
        # TODO: need a decorator to capture this output and put it into the doc,
        #       or maybe just a new kind of assert?

    def test_feature_run_driver(self):
        prob = Problem()
        model = prob.model = SellarDerivatives()
        model.nonlinear_solver = NonlinearBlockGS()

        prob.driver = ScipyOptimizer()
        prob.driver.options['optimizer'] = 'SLSQP'
        prob.driver.options['tol'] = 1e-9

        model.add_design_var('z', lower=np.array([-10.0, 0.0]), upper=np.array([10.0, 10.0]))
        model.add_design_var('x', lower=0.0, upper=10.0)
        model.add_objective('obj')
        model.add_constraint('con1', upper=0.0)
        model.add_constraint('con2', upper=0.0)

        prob.setup()
        prob.run_driver()

        assert_rel_error(self, prob['x'], 0.0, 1e-5)
        assert_rel_error(self, prob['y1'], 3.160000, 1e-2)
        assert_rel_error(self, prob['y2'], 3.755278, 1e-2)
        assert_rel_error(self, prob['z'], [1.977639, 0.000000], 1e-2)
        assert_rel_error(self, prob['obj'], 3.18339395, 1e-2)

    def test_feature_promoted_sellar_set_get_outputs(self):

        prob = Problem()
        prob.model = SellarDerivatives()
        prob.model.nonlinear_solver = NonlinearBlockGS()

        prob.setup()

        prob['x'] = 2.75
        assert_rel_error(self, prob['x'], 2.75, 1e-6)

        prob.run_model()

        assert_rel_error(self, prob['y1'], 27.3049178437, 1e-6)

    def test_feature_not_promoted_sellar_set_get_outputs(self):

        prob = Problem()
        prob.model = SellarDerivativesConnected()
        prob.model.nonlinear_solver = NonlinearBlockGS()

        prob.setup()

        prob['px.x'] = 2.75
        assert_rel_error(self, prob['px.x'], 2.75, 1e-6)

        prob.run_model()

        assert_rel_error(self, prob['d1.y1'], 27.3049178437, 1e-6)

    def test_feature_promoted_sellar_set_get_inputs(self):

        prob = Problem()
        prob.model = SellarDerivatives()
        prob.model.nonlinear_solver = NonlinearBlockGS()

        prob.setup()

        prob['x'] = 2.75
        assert_rel_error(self, prob['x'], 2.75, 1e-6)

        prob.run_model()

        # the output variable, referenced by the promoted name
        assert_rel_error(self, prob['y1'], 27.3049178437, 1e-6)
        # the connected input variable, referenced by the absolute path
        assert_rel_error(self, prob['d2.y1'], 27.3049178437, 1e-6)

    def test_feature_set_get_array(self):
        prob = Problem()
        prob.model = SellarDerivatives()
        prob.model.nonlinear_solver = NonlinearBlockGS()

        prob.setup()

        # default value from the class definition
        assert_rel_error(self, prob['x'], 1.0, 1e-6)
        prob['x'] = 2.75
        assert_rel_error(self, prob['x'], 2.75, 1e-6)

        assert_rel_error(self, prob['z'], [5.0, 2.0], 1e-6)
        prob['z'] = [1.5, 1.5]  # for convenience we convert the list to an array.
        assert_rel_error(self, prob['z'], [1.5, 1.5], 1e-6)
        prob['z'] = [1.5, 1.5]  # for convenience we convert the list to an array.
        assert_rel_error(self, prob['z'], (1.5, 1.5), 1e-6)

        prob.run_model()
        assert_rel_error(self, prob['y1'], 5.43379016853, 1e-6)
        assert_rel_error(self, prob['y2'], 5.33104915618, 1e-6)

        prob['z'] = np.array([2.5, 2.5])
        assert_rel_error(self, prob['z'], [2.5, 2.5], 1e-6)

        prob.run_model()
        assert_rel_error(self, prob['y1'], 9.87161739688, 1e-6)
        assert_rel_error(self, prob['y2'], 8.14191301549, 1e-6)

    @unittest.skip('access via promoted names is not working yet')
    def test_feature_residuals(self):

        prob = Problem()
        prob.model = SellarDerivatives()
        prob.model.nonlinear_solver = NonlinearBlockGS()

        prob.setup()

        # default value from the class definition

        prob['z'] = [1.5, 1.5]  # for convenience we convert the list to an array.
        prob.run_model()

        inputs, outputs, residuals = prob.model.get_nonlinear_vectors()

        self.assertLess(residuals['y1'], 1e-6)
        self.assertLess(residuals['y2'], 1e-6)

    def test_setup_bad_mode(self):
        # Test error message when passing bad mode to setup.

        prob = Problem()

        try:
            prob.setup(mode='junk')
        except ValueError as err:
            msg = "Unsupported mode: 'junk'"
            self.assertEqual(str(err), msg)
        else:
            self.fail('Expecting ValueError')

    def test_run_before_setup(self):
        # Test error message when running before setup.

        prob = Problem()

        try:
            prob.run_model()
        except RuntimeError as err:
            msg = "The `setup` method must be called before `run_model`."
            self.assertEqual(str(err), msg)
        else:
            self.fail('Expecting RuntimeError')

        try:
            prob.run_driver()
        except RuntimeError as err:
            msg = "The `setup` method must be called before `run_driver`."
            self.assertEqual(str(err), msg)
        else:
            self.fail('Expecting RuntimeError')

    def test_root_deprecated(self):
        # testing the root property
        msg = "The 'root' property provides backwards compatibility " \
            + "with OpenMDAO <= 1.x ; use 'model' instead."

        prob = Problem()

        # check deprecation on setter
        with warnings.catch_warnings(record=True) as w:
            prob.root = Group()

        self.assertEqual(len(w), 1)
        self.assertTrue(issubclass(w[0].category, DeprecationWarning))
        self.assertEqual(str(w[0].message), msg)

        # check deprecation on getter
        with warnings.catch_warnings(record=True) as w:
            prob.root

        self.assertEqual(len(w), 1)
        self.assertTrue(issubclass(w[0].category, DeprecationWarning))
        self.assertEqual(str(w[0].message), msg)

        # testing the root kwarg
        with self.assertRaises(ValueError) as cm:
            prob = Problem(root=Group(), model=Group())
        err = cm.exception
        self.assertEqual(str(err), "cannot specify both `root` and `model`. `root` has been "
                         "deprecated, please use model")

        with warnings.catch_warnings(record=True) as w:
            prob = Problem(root=Group)

        self.assertEqual(str(w[0].message), "The 'root' argument provides backwards "
                         "compatibility with OpenMDAO <= 1.x ; use 'model' instead.")

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