import numpy as np
import scipy.signal
import scipy.interpolate

from .filterfactory import (gammatone, asymmetric_compensation_coeffs, loggammachirp, erbspace)
from .tools import sosfreqz

NFFT = 2**16

def make_cgc_filter(freqs, compression_degree, level_max, level_step, sample_rate, dtype='float32'):
    """

    parameters
    ----

    freqs: vector of central freqquencies Hz

    compression_degree: vector of compression degree for each freq with:
        * 1=no compression impairement
        * 0= maximum compression impairement

    returns
    ---

        coefficients_pgc: sos coefficient of pgc filters shape is (freqs.size, 8, 6)

        coefficients_hpaf: sos coefficient of hpaf (level dependant). Filter shape is
            (freqs.size, levels.size, 4, 6)

        levels: vector of levels.

        band_overlap_gain: gain for conpensating the overlap between bands. this depend
          the number of band

    """
    freqs = np.asarray(freqs)
    compression_degree = np.asarray(compression_degree)

    nb_freq_band = len(freqs)

    # pgc filter coefficient
    b1 = 1.81
    c1 = -2.96

    # hpaf filter coefficient
    b2 = 2.17
    c2 = 2.2

    p0=2
    p1=1.7818*(1-0.0791*b2)*(1-0.1655*abs(c2))
    p2=0.5689*(1-0.1620*b2)*(1-0.0857*abs(c2))
    p3=0.2523*(1-0.0244*b2)*(1+0.0574*abs(c2))
    p4=1.0724

    coefficients_pgc = loggammachirp(freqs, sample_rate, b=b1, c=c1).astype(dtype)

    #noramlize PGC to 0 db at maximum
    #pcg_freqs = np.zeros_like(freqs)  #for testing
    for f, freq in enumerate(freqs):
        w, h = sosfreqz(coefficients_pgc[f,:,:], worN =2**16,)
        gain = np.max(np.abs(h))
        coefficients_pgc[f,0, :3] /= gain
        #pcg_freqs[f] = (w/np.pi*(sample_rate/2.))[np.argmax(np.abs(h))] #for testing

    # Construct hpaf filters : pre compute for all sound levels for each freq
    levels = np.arange(0, level_max+level_step,level_step)
    nlevel = levels.size

    # construct hpaf depending on compression_degree for each freq
    coefficients_hpaf = np.zeros((nb_freq_band, len(levels), 4, 6), dtype = dtype)

    alpha = compression_degree

    # <NAME> coefficient (if this is correct)
    # need to be check with irino
    frat0r = 1 + 0.466 * (1-alpha)
    frat1r =  - 0.0109 * (1-alpha)

    # <NAME> coefficient 2015
    #~ w = (1-alpha) * 5 / 3
    #~ frat1r = -w/65/2.
    #~ frat0r = 1+w/2.

    # <NAME> coefficient 2016
    #~ frat0r = 1 + (1-alpha)*1.3333
    #~ frat1r = - (1-alpha) * 0.0205

    for l, level in enumerate(levels):
        # minus for inverse compression = moving left
        frat = frat0r + frat1r * level
        freqs2 = freqs*frat
        #freqs2 = pcg_freqs*frat #for testing
        coefficients_hpaf[:, l, : , : ] = asymmetric_compensation_coeffs(freqs2, sample_rate, b2,c2,p0,p1,p2,p3,p4)

    #noramlize for highest level
    for f, freq in enumerate(freqs):
        #~ print('freq', freq)
        filter = np.concatenate([coefficients_pgc[f,:,:], coefficients_hpaf[f , -1, : , : ],coefficients_pgc[f,:,:] ], axis = 0)
        w, h = sosfreqz(filter, worN =NFFT)
        gain = np.max(np.abs(h))
        coefficients_hpaf[f , :, 0 , :3 ] /= gain

    # compensate final gain of sum of all band freqs
    all = np.zeros(NFFT)
    for f, freq in enumerate(freqs):
        #~ print('freq', freq)
        all_filter = np.concatenate([coefficients_pgc[f,:,:],coefficients_hpaf[f,-1,:,:], coefficients_pgc[f,:,:]], axis = 0)
        w, h = sosfreqz(all_filter,worN = NFFT)
        all += np.abs(h)

    # check this
    fft_freqs = w/np.pi*(sample_rate/2.)
    #all = all[(fft_freqs>freqs[0]) & (fft_freqs<freqs[-1])]

    # this is bad because of side effect
    #dbgain_final = -np.mean(20*np.log10(all))

    band_overlap_gain_db = -np.max(20*np.log10(all))

    band_overlap_gain = 10**(band_overlap_gain_db/20.)

    #~ print('band_overlap_gain_db', band_overlap_gain_db)

    return coefficients_pgc, coefficients_hpaf, levels, band_overlap_gain

def make_invcomp_filter(freqs, compression_degree, level_max, level_step, sample_rate, dtype='float32'):
    """
    Similar than make_cgc_filter except coefficients_hpaf is replace by gain_controlled.

    Used by the class InvComp.

    parameters
    ----

    freqs: vector of central freqquencies Hz

    compression_degree: vector of compression degree for each freq with:
        * 1=no compression impairement
        * 0= maximum compression impairement

    returns
    ---

        coefficients_pgc: sos coefficient of pgc filters shape is (freqs.size, 8, 6)

        gain_controlled: gain level dependant. Filter shape is (freqs.size, levels.size)

        levels: vector of levels.

        band_overlap_gain: gain for conpensating the overlap between bands. this depend
          the number of band

    Note that  for computing gain_controlled we use the same dynamic that InvCGC.
    So internally HPAF is computed to extract the equivalent gain level controlled.
    Also note than this HPAF is alwas done at 1000Hz because the gain response is
    better in the middle of the spectrum.

    """
    freqs = np.asarray(freqs)
    compression_degree = np.asarray(compression_degree)

    nb_freq_band = len(freqs)

    # pgc filter coefficient
    b1 = 1.81
    c1 = -2.96

    # hpaf filter coefficient
    b2 = 2.17
    c2 = 2.2

    p0=2
    p1=1.7818*(1-0.0791*b2)*(1-0.1655*abs(c2))
    p2=0.5689*(1-0.1620*b2)*(1-0.0857*abs(c2))
    p3=0.2523*(1-0.0244*b2)*(1+0.0574*abs(c2))
    p4=1.0724

    coefficients_pgc = loggammachirp(freqs, sample_rate, b=b1, c=c1).astype(dtype)

    #noramlize PGC to 0 db at maximum
    #pcg_freqs = np.zeros_like(freqs)  #for testing
    for f, freq in enumerate(freqs):
        w, h = sosfreqz(coefficients_pgc[f,:,:], worN =2**16,)
        gain = np.max(np.abs(h))
        coefficients_pgc[f,0, :3] /= gain
        #pcg_freqs[f] = (w/np.pi*(sample_rate/2.))[np.argmax(np.abs(h))] #for testing

    # Construct hpaf filters : pre compute for all sound levels for each freq
    levels = np.arange(0, level_max+level_step,level_step)
    nlevel = levels.size

    # construct hpaf depending on compression_degree for each freq

    #for efficieny we compute only once each compression_degree for deifferent bands
    _freqs = np.array([1000.])
    pgc_1000Hz = loggammachirp(_freqs, sample_rate, b=b1, c=c1).astype(dtype)
    pgc_1000Hz = pgc_1000Hz[0,:,:]

    w, h = sosfreqz(pgc_1000Hz, worN =2**16,)
    gain = np.max(np.abs(h))
    pgc_1000Hz[0, :3] /= gain
    hpaf_1000Hz = np.zeros((len(levels), 4, 6), dtype = dtype)
    gain_controlled_by_alpha = {}
    #~ print('all alpha', np.unique(compression_degree))
    for alpha in np.unique(compression_degree):
        print('alpha', alpha)
        # Toshio Irino coefficient (if this is correct)
        # need to be check with irino
        frat0r = 1 + 0.466 * (1-alpha)
        frat1r =  - 0.0109 * (1-alpha)

        for l, level in enumerate(levels):
            print('level', level)
            # minus for inverse compression = moving left
            frat = frat0r + frat1r * level
            freqs2 = _freqs*frat
            #freqs2 = pcg_freqs*frat #for testing
            hpaf_1000Hz[l,:,:] = asymmetric_compensation_coeffs(freqs2, sample_rate, b2,c2,p0,p1,p2,p3,p4)

        #noramlize for highest level
        filter = np.concatenate([pgc_1000Hz, hpaf_1000Hz[-1, : , : ],pgc_1000Hz], axis = 0)
        w, h = sosfreqz(filter, worN =NFFT)
        gain = np.max(np.abs(h))
        hpaf_1000Hz[:, 0 , :3 ] /= gain

        gains = np.zeros(len(levels))
        for l, level in enumerate(levels):
            print('level', level)
            filter = np.concatenate([pgc_1000Hz, hpaf_1000Hz[l, : , : ],pgc_1000Hz], axis = 0)
            w, h = sosfreqz(filter, worN =NFFT)
            #~ gains[l] = np.max(20*np.log10(np.abs(h)))
            gains[l] = np.max(np.abs(h))
        gain_controlled_by_alpha[alpha] = gains

    #~ alpha = compression_degree

    # <NAME> coefficient 2015
    #~ w = (1-alpha) * 5 / 3
    #~ frat1r = -w/65/2.
    #~ frat0r = 1+w/2.

    # <NAME> coefficient 2016
    #~ frat0r = 1 + (1-alpha)*1.3333
    #~ frat1r = - (1-alpha) * 0.0205

    #~ for l, level in enumerate(levels):
        # minus for inverse compression = moving left
        #~ frat = frat0r + frat1r * level
        #~ freqs2 = freqs*frat
        #freqs2 = pcg_freqs*frat #for testing
        #~ coefficients_hpaf[:, l, : , : ] = asymmetric_compensation_coeffs(freqs2, sample_rate, b2,c2,p0,p1,p2,p3,p4)

    #noramlize for highest level
    #~ for f, freq in enumerate(freqs):
        #~ filter = np.concatenate([coefficients_pgc[f,:,:], coefficients_hpaf[f , -1, : , : ],coefficients_pgc[f,:,:] ], axis = 0)
        #~ w, h = sosfreqz(filter, worN =NFFT)
        #~ gain = np.max(np.abs(h))
        #~ coefficients_hpaf[f , :, 0 , :3 ] /= gain

    gain_controlled = np.zeros((freqs.size, levels.size), dtype=dtype)
    for f, freq in enumerate(freqs):
        alpha = compression_degree[f]
        gain_controlled[f, :] = gain_controlled_by_alpha[alpha]

    # compensate final gain of sum of all band freqs
    all = np.zeros(NFFT)
    for f, freq in enumerate(freqs):
        all_filter = np.concatenate([coefficients_pgc[f,:,:], coefficients_pgc[f,:,:]], axis = 0)
        w, h = sosfreqz(all_filter,worN = NFFT)
        all += np.abs(h)

    # check this
    fft_freqs = w/np.pi*(sample_rate/2.)
    #all = all[(fft_freqs>freqs[0]) & (fft_freqs<freqs[-1])]

    # this is bad because of side effect
    #dbgain_final = -np.mean(20*np.log10(all))

    band_overlap_gain_db = -np.max(20*np.log10(all))

    band_overlap_gain = 10**(band_overlap_gain_db/20.)

    #~ print('band_overlap_gain_db', band_overlap_gain_db)

    return coefficients_pgc, gain_controlled, levels, band_overlap_gain


#
# Description: example of a bootstrapping and forward curve generation
# script, this can be used to build a set of curves for different currencies
# TODO: include some spline smoothing to the zero curve, from first principles!
#

from sympy.solvers import solve
from sympy import Symbol, abs, Real
x = Symbol('x', real=True)

import pylab as pylab

def g(yieldCurve, zeroRates,n, verbose):
    '''
        generates recursively the zero curve
        expressions eval('(0.06/1.05)+(1.06/(1+x)**2)-1')
        solves these expressions to get the new rate
        for that period

    '''
    if len(zeroRates) >= len(yieldCurve):
        print("\n\n\t+zero curve boot strapped [%d iterations]" % (n))
        return
    else:
        legn = ''
        for i in range(0,len(zeroRates),1):
            if i == 0:
                legn = '%2.6f/(1+%2.6f)**%d'%(yieldCurve[n], zeroRates[i],i+1)
            else:
                legn = legn + ' +%2.6f/(1+%2.6f)**%d'%(yieldCurve[n], zeroRates[i],i+1)
        legn = legn + '+ (1+%2.6f)/(1+x)**%d-1'%(yieldCurve[n], n+1)
        # solve the expression for this iteration
        if verbose:
            print("-[%d] %s" % (n, legn.strip()))
        rate1 = solve(eval(legn), x)
        # Abs here since some solutions can be complex
        rate1 = min([Real(abs(r)) for r in rate1])
        if verbose:
            print("-[%d] solution %2.6f" % (n, float(rate1)))
        # stuff the new rate in the results, will be
        # used by the next iteration
        zeroRates.append(rate1)
        g(yieldCurve, zeroRates,n+1, verbose)

verbose = True
tenors = [.1,.25,0.5,1,2,3,5,7,10,20,30]
#
# money market, futures, swap rates
#
yieldCurve = [0.07,	0.09,	0.15,	0.21,	0.37,	0.57,	1.13,	1.70,	2.31,	3.08	,3.41]

#yieldCurve = [0.05, 0.06, 0.07, 0.08 ,0.085 ,0.0857 ,0.0901,0.0915,0.0925,0.0926,0.0934,0.0937]
zeroRates = [yieldCurve[0]] # TODO: check that this is the correct rate

print("\n\n\t<NAME>, March 2012\n\tYield Curve Bootstrapper\n\t<NAME>\n\n")

# kick off the recursive code
g(yieldCurve, zeroRates, 1, verbose)
print("\tZeroRate Array",zeroRates)

pylab.plot(tenors,yieldCurve)
pylab.plot(tenors,zeroRates)
pylab.show()


<gh_stars>1-10
import torch
import cv2
import numpy as np

from skimage.transform import PiecewiseAffineTransform, warp
from scipy.ndimage.interpolation import map_coordinates
from scipy.ndimage.filters import gaussian_filter
from skimage import io, transform
from PIL import Image

class Rescale(object):
    """Rescale the image in a sample to a given size.

    Args:
        output_size (tuple or tuple): Desired output size. If tuple, output is
            matched to output_size. If int, smaller of image edges is matched
            to output_size keeping aspect ratio the same.
    """

    def __init__(self, output_size, keep_ratio=False):
        self.output_size = output_size
        self.keep_ratio = keep_ratio

    def __call__(self, sample):
        image = sample
        image = np.array(image)

        h, w = image.shape[:2]
        if isinstance(self.output_size, int):
            new_h, new_w = self.output_size, self.output_size * w / h
            new_h, new_w = int(new_h), int(new_w)
            img = cv2.resize(image,(new_w, new_h))
        else:
            if self.keep_ratio:
                ratio = w/h
                img_h, img_w = self.output_size
                new_h = img_h
                new_w = min(img_w, new_h * ratio)
                new_h, new_w = int(new_h), int(new_w)
                img = cv2.resize(image,(new_w, new_h))
                pad = img_w - new_w
                # img = np.pad(img, ((0, 0),(0, pad)), 'maximum')
                img = np.pad(img, ((0, 0),(0, pad)), 'constant', constant_values=(255, ))
            else:
                new_h, new_w = self.output_size
                new_h, new_w = int(new_h), int(new_w)
                img = cv2.resize(image,(new_w, new_h))
        # print(img.shape)
        # img = Image.fromarray(img)
        return img

class ElasticTransformation(object):
    """ElasticTransformation the image in a sample to a given size.
    Code adapted from https://www.kaggle.com/bguberfain/elastic-transform-for-data-augmentation

    Reducing sigma increases distortion
    increasing alpha increases distortion
    """
    def __init__(self, prob, alpha=0.3, sigma=0.04, borderValue=255):
        self.prob = prob
        self.borderValue = borderValue
        self.sigma = sigma
        self.alpha = alpha

    def __call__(self, sample):

        image = np.array(sample)

        if(np.random.rand()>self.prob):
            return image
        alpha = image.shape[0]*self.alpha
        sigma = image.shape[0]*self.sigma
        random_state = None

        if random_state is None:
            random_state = np.random.RandomState(None)

        shape = image.shape
        shape_size = shape[:2]
        dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode='reflect',
                                         cval=self.borderValue) * alpha
        dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode='reflect',
                                         cval=self.borderValue) * alpha

        x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
        indices = np.reshape(y+dy, (-1, 1)), np.reshape(x+dx, (-1, 1))

        newimage = map_coordinates(image, indices, order=1, mode='constant',
                                         cval=self.borderValue).reshape(shape)
        return newimage

class AffineTransformation(object):
    """AffineTransformation
    """
    def __init__(self, prob, rotate=2, shear=0.5, borderValue=255):
        self.prob = prob
        self.borderValue = borderValue
        self.rotate = rotate
        self.shear = shear
        assert shear <= 1, 'Max shear value is 1'

    def __call__(self, sample):
        image = np.array(sample)

        if(np.random.rand()>self.prob):
            return Image.fromarray(image)

        rows, cols = image.shape
        if(np.random.rand()<0.5):
            rotAngle = np.random.randint(-self.rotate, self.rotate)

            height, width = image.shape[:2]
            image_center = (width/2, height/2)

            rotation_mat = cv2.getRotationMatrix2D(image_center, rotAngle, 1.)

            abs_cos = abs(rotation_mat[0,0])
            abs_sin = abs(rotation_mat[0,1])

            bound_w = int(height * abs_sin + width * abs_cos)
            bound_h = int(height * abs_cos + width * abs_sin)

            rotation_mat[0, 2] += bound_w/2 - image_center[0]
            rotation_mat[1, 2] += bound_h/2 - image_center[1]

            image = cv2.warpAffine(image, rotation_mat, (bound_w, bound_h), borderMode=cv2.BORDER_CONSTANT, borderValue=self.borderValue)

        if(np.random.rand()<0.5):
            shearAngle = -self.shear + np.random.rand()
            M = np.array([[1.0,shearAngle,0.0],[0.0,1.0,0.0]])
            image = cv2.warpAffine(image, M, (cols,rows), borderMode=cv2.BORDER_CONSTANT, borderValue=self.borderValue)
        image = cv2.resize(image,(cols,rows))
        return Image.fromarray(image)

class MultiScale(object):
    def __init__(self, prob, scale_axis=2, scale_factor=2):
        self.prob = prob
        self.scale_axis = scale_axis
        self.scale_factor = scale_factor

    def __call__(self, sample):
        image = sample
        h, w = image.shape
        if np.random.rand()>self.prob:
            return image

        colors_ = np.unique(image)
        if np.max(colors_) >= 170:
            color = colors_[colors_>=170].mean()
        else:
            color = 255
        color = 255
        nw = w
        nh = h
        if self.scale_axis == 0 or self.scale_axis == 2:
            nw = w//(2**np.random.randint(1, self.scale_factor))
        if self.scale_axis == 1 or self.scale_axis == 2:
            nh = h//(2**np.random.randint(1, self.scale_factor))

        rimage = cv2.resize(image, (nw, nh))
        image = color*np.ones((h, w))
        # print(image.shape, rimage.shape)
        if self.scale_axis == 2:
            nx = np.random.randint(3, w-nw)
            ny = np.random.randint(3, h-nh)
            image[ny:ny+nh, nx:nx+nw] = rimage
        elif self.scale_axis == 0:
            nx = np.random.randint(3, w-nw)
            image[0:h, nx:nx+nw] = rimage
        elif self.scale_axis == 1:
            ny = np.random.randint(3, h-nh)
            image[ny:ny+nh, 0:w] = rimage
        return image

class PiecewiseAffine(object):
    def __init__(self, prob, dfactor=10, pfactor=(0, 2)):
        self.prob = prob
        self.dfactor = dfactor
        self.pfactor = pfactor

    def __call__(self, sample):
        image = sample
        if np.random.rand()>self.prob:
            return image

        image = np.array(image)
        rows, cols = image.shape[0], image.shape[1]
        src_cols = np.linspace(0, cols, 4)
        src_rows = np.linspace(0, rows, 2)
        src_rows, src_cols = np.meshgrid(src_rows, src_cols)
        src = np.dstack([src_cols.flat, src_rows.flat])[0]

        dfactor = np.random.randint(10, 20)
        pfactor = (np.random.randint(0, 3), np.random.randint(2, 4))
        dst_rows = src[:, 1] - np.sin(np.linspace(pfactor[0]*np.pi/2, pfactor[1]*np.pi, src.shape[0]))*dfactor
        dst_cols = src[:, 0]
        dst = np.vstack([dst_cols, dst_rows]).T

        tform = PiecewiseAffineTransform()
        tform.estimate(src, dst)
        out_image = warp(image, tform, output_shape=(rows, cols), cval=255, preserve_range=True)
        out_image = Image.fromarray(np.uint8(out_image))
        return out_image

class Normalize(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample):
        image = sample
        image = (image-np.mean(image)) / ((np.std(image) + 0.0001) / 128.0)

        return image

class ColorInvert(object):
    """Convert ndarrays in sample to Tensors."""
    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, sample):
        image = sample

        if np.random.rand()>self.prob:
            return image

        cnst = 255
        image = cnst - image
        return image

class ToTensor(object):
    """Convert arrays in sample to Tensors."""

    def __call__(self, sample):
        image = np.array(sample)
        image = torch.from_numpy(np.array(image))
        image = image.float().sub(128).div(128)
        image = image.unsqueeze(0)
        return image

class NormalizedTensor(object):
    def __call__(self, sample):
        image = np.array(sample)
        image = torch.from_numpy(np.array(image, dtype=np.double))
        image.sub_(0.5).div_(0.5)
        image = image.unsqueeze(0)
        return image


import numpy as np
import datetime, sys
import scipy.io
import pdb
import os
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from shapely.geometry import Polygon,Point
import scipy.interpolate
import xarray as xr
import matplotlib.pyplot as plt
import utilities.hfr_util
import pyproj
from matplotlib import animation, rc
plt.rcParams['animation.html'] = 'jshtml'
animation.rcParams['animation.embed_limit'] = 60

class TrackingModel(object):

    def __init__(self):
        self.minlon, self.maxlon, self.minlat, self.maxlat, self.resolution_km, self.url, self.days_to_capture, self.start_time = self.set_model_parameters(default=True)
        self.current_dataset = self._get_HFR_subset()
        self.lat = self.current_dataset['lat'].values
        self.lon = self.current_dataset['lon'].values
        self.x = np.arange(0, len(self.lon), 1) * 1850
        self.y = np.arange(0, len(self.lat), 1) * 1995
        self.x_grid, self.y_grid = np.meshgrid(self.x, self.y)
        self.lon_grid, self.lat_grid = np.meshgrid(self.lon, self.lat)
        self.origin = ccrs.UTM(zone=10).transform_point(self.lon[0], self.lat[0], ccrs.PlateCarree())
        self.time_index = 0
        self.hours_elapsed = 0
        self.current_time = self.start_time
        self.particles = np.array([])
        self.time_step = .25

    def set_model_parameters(self, default=False):
        ''' Hotwire the default parameters for developement '''
        crs = ccrs.PlateCarree() # Make cartopy projection object

        resolution_km = 2
        days_to_capture=3

        if int(resolution_km)==2:
            url='http://hfrnet-tds.ucsd.edu/thredds/dodsC/HFR/USWC/2km/hourly/RTV/HFRADAR_US_West_Coast_2km_Resolution_Hourly_RTV_best.ncd'
        if int(resolution_km)==6:
            url='http://hfrnet-tds.ucsd.edu/thredds/dodsC/HFR/USWC/6km/hourly/RTV/HFRADAR_US_West_Coast_6km_Resolution_Hourly_RTV_best.ncd'

        if default:
            year,month,day,hour,minu,secs = 2017,5,27,5,0,0
            start_time = datetime.datetime(year,month,day,hour,minu,secs)
            maxlat, minlat = 38.5, 37
            maxlat=float(maxlat)
            minlat=float(minlat)
            minlon, maxlon = -123.5, -122
            minlon=float(minlon)
            maxlon=float(maxlon)

        else:
            try:
                year,month,day,hour,minu,secs=input("Enter the start date [yr mn dy hr mn sc] ").split()
            except ValueError:
                print('Defaulting to 2017-05-27T05:00:00')
                year,month,day,hour,minu,secs = 2017,5,27,5,0,0
                start_time = datetime.datetime(year,month,day,hour,minu,secs)
            #[year,month,day,hour,minu,secs]=input('Enter the start date [yr,mn,dy,hr,mn,sc]): ')
            #[2017,5,27,5,0,0]
            try:
                maxlat,minlat=input('Enter the north and south boundary []: ').split()
            except ValueError:
                maxlat, minlat = 38.5, 37
            maxlat=float(maxlat)
            minlat=float(minlat)
            # [38.5,37]
            try:
                minlon, maxlon=input('Enter the west and east boundary []: ').split()
            except ValueError:
                minlon, maxlon = -123.5, -122
            minlon=float(minlon)
            maxlon=float(maxlon)

        return minlon, maxlon, minlat, maxlat, resolution_km, url, days_to_capture, start_time

    def _get_HFR_subset(self):

        """
        Open netcdf from hfrnet as xarray datase, subset the desired data in space and time
        Currently, using global variables for parameters, should reduce as code becomes object-oriented

        :return hfr_ds = xarray dataset over the specified spatial grid and time
        """
        try:
            ds = xr.open_dataset(self.url)
        except Exception as e:
            print("Trouble downloading data, check connections.")
            print(e)
            sys.exit()

        subset_ds = ds.sel(time = slice(self.start_time, self.start_time + datetime.timedelta(days=self.days_to_capture)),
                      lat =  slice(self.minlat,self.maxlat),
                      lon = slice(self.minlon,self.maxlon))
        clean_u = hfr_util.interp_time_surface_currents(subset_ds['u'].values) # linear interpolate currents through time, given a threshold of availible data
        clean_v = hfr_util.interp_time_surface_currents(subset_ds['v'].values)
        subset_ds['u_clean'] =  xr.DataArray(clean_u, coords={'time': subset_ds['time'].values, 'lon': subset_ds['lon'].values, 'lat': subset_ds['lat'].values}, dims=['time','lat','lon'])
        subset_ds['v_clean'] =  xr.DataArray(clean_v, coords={'time': subset_ds['time'].values, 'lon': subset_ds['lon'].values, 'lat': subset_ds['lat'].values}, dims=['time','lat','lon'])
        return subset_ds

    def add_particle(self, coord):
        utm = ccrs.UTM(zone=10)
        p_dist = utm.transform_point(coord[0], coord[1], ccrs.PlateCarree())
        p_dist = [p_dist[0] - self.origin[0], p_dist[1] - self.origin[1]]
        self.particles = np.append(self.particles, Particle(p_dist))

    def part_to_coor(self, particle, last=False):
        pc_proj = ccrs.PlateCarree()
        p_dist = particle.coordinates
        p_dist = [p_dist[:,0] + self.origin[0], p_dist[:,1] + self.origin[1]]
        if last:
            p_coor = pc_proj.transform_points(ccrs.UTM(zone=10), np.array((p_dist[:,-1])),np.array((p_dist[:,-1])))
        else:
            p_coor = pc_proj.transform_points(ccrs.UTM(zone=10), np.array([p_dist[:,0]]),np.array([p_dist[:,1]]))
        p_coor = p_coor[:,:2]
        return p_coor

    def draw_map(self, draw_bathy=True, whole_extent=False):
        ''' Draw a map of the domain '''

        land_feature  = cfeature.NaturalEarthFeature('physical','land','10m')
        self.fig = plt.figure()
#        self.fig.set_size_inches(8,8)
        self.geo_axes = plt.axes(projection=ccrs.PlateCarree())
        if whole_extent:
            extent = [self.minlon, self.maxlon, self.minlat, self.maxlat]
        else:
            extent = [-123, -122.3, 37.26, 38]
        self.geo_axes.set_extent(extent, crs=ccrs.PlateCarree())
        self.geo_axes.add_feature(land_feature, edgecolor='k', zorder=40)
#        self.geo_axes.gridlines(draw_labels=True, zorder= 20)
        if draw_bathy:
            self._draw_bathymetry_SFBAY()

    def _draw_bathymetry_SFBAY(self):
        '''
        Draw and return an axes with bathymetry contours

        Bathy grid retrieved from: https://maps.ngdc.noaa.gov/viewers/wcs-client/
        '''
        file = './data/SF_Bay_coastal_relief_model.nc'
        bathy = xr.open_dataset(file)
        blat = bathy['lat'].values
        blon = bathy['lon'].values
        elev = bathy['Band1'].values
        levels = np.arange(10*(-2000//10), -499, 450)
        levels = np.append(levels, np.arange(-475, 25, 25))
        lws = [0.5 if l%100 else 1 for l in levels]
        cs = self.geo_axes.contour(blon, blat, elev, levels, linewidths=lws, linestyles='solid', colors=['black'], alpha=0.4)
        plt.clabel(cs, list(np.arange(-125,1,25)),fmt='%1d', inline=True, fontsize=15, colors='k',inline_spacing=10)

    def plot_particles(self, plot_type='all', save_dir):
        """
        Make plots of the particle trajectories. Can make static plots, last locations, or animations
        """
        if plot_type == 'last':
            #Only if you want to plot the last position (debugging)
            for p in self.particles[:]:
                pos = self.part_to_coor(p, last=True)
                self.geo_axes.scatter(pos[0,0],pos[0,1], zorder=50, marker='.', c='grey')

        elif plot_type == 'animation':
            # make animation of all the particle trajectories
            Q = self.geo_axes.quiver(self.lon_grid, self.lat_grid, model.current_dataset['u_clean'].isel(time=0), model.current_dataset['v_clean'].isel(time=0), color='b', units='inches')
            p = self.particles[0]
            px, py = p.coordinates[:,0]+self.origin[0], p.coordinates[:,1]+self.origin[1]
            line, = self.geo_axes.plot(px[0],py[0],
                                       color='.75',
                                       transform=ccrs.UTM(zone=10))
            front_marker = self.geo_axes.scatter(px[0],py[0],
                                                 color='b',
                                                 s=10,
                                                 transform=ccrs.UTM(zone=10))
            tail_marker = self.geo_axes.scatter(px[0],py[0],
                                                marker='x',
                                                color='b',
                                                s=10,
                                                transform=ccrs.UTM(zone=10)) # last
            line, = self.geo_axes.plot(px[0],py[0],
                                       color='.75',
                                       transform=ccrs.UTM(zone=10))
            time_text = self.geo_axes.text(0.5, 1.05, 'Hours Before: ' + str(0), fontsize=16, transform=self.geo_axes.transAxes,zorder=300)

            def update_geoaxes(num, front_marker, tail_marker, line, Q, time_text):
                """
                Animation function - updates the data of each plot at each
                timestep and the hour textbox
                """
                Q.set_UVC(model.current_dataset['u_clean'].isel(time=num//4),model.current_dataset['v_clean'].isel(time=num//4))
                front_marker.set_offsets((px[num],py[num]))
                if num == 0:
                    tail = 0
                    line.set_data(px[num], py[num])

                elif num <= 24:
                    tail = 0
                    line.set_data(px[tail:num], py[tail: num])
                else:
                    tail = num - 24
                    line.set_data(px[tail:num], py[tail: num])

                tail_marker.set_offsets((px[tail], py[tail]))
                time_text.set_text('Hours Before: ' + str(num//4))

                return(front_marker, tail_marker, line, Q, time_text)

            anim = animation.FuncAnimation(self.fig,
                                           update_geoaxes, fargs=(front_marker, tail_marker, line, Q, time_text),
                                           interval=150,
                                           blit=True,
                                           frames=len(px))
            Writer = animation.writers['ffmpeg']
            writer = Writer(fps=15, metadata=dict(artist='<NAME>'), bitrate=1800)
            anim.save('testing_traj.mp4', dpi=200, writer=writer)

        elif plot_type == 'all':
            for p in self.particles[:]:
                self.geo_axes.plot(p.coordinates[:,0]+self.origin[0], p.coordinates[:,1]+self.origin[1],
                                   zorder=50,
                                   c='blue',
                                   markerfacecolor='.5',
                                   markeredgecolor='None',
                                   marker='.',
                                   markersize=5,
                                   transform=ccrs.UTM(zone=10))
        else:
             raise ValueError("Plot type not valid.")

    def advect_particle(self):
        """
        Advect a particle object based on the current
        Fixed timestep to 1/4 hour --> If this is to be worked update self.timestep
        """
        for p in self.particles:
            current_pos = p.get_position()
            if not np.any(np.isnan(current_pos)): # check if advection previously failed if so skip
                new_x = self._solve_position(current_pos[0], current_pos[1], self.time_step, dim=0)
                new_y = self._solve_position(current_pos[0], current_pos[1], self.time_step, dim=1)
                p.update_position([new_x, new_y])
        self.update_time()

    def update_time(self):
        """
        Update the time_index and number of elapsed hours with each advection
        """

        self.hours_elapsed += self.time_step
        self.time_index = int(self.hours_elapsed)

    def _solve_position(self, x, y, h, dim):
            """
            Solves for the next position of a particle in after time, h, in either the x
            or y using a runge-kutta 4th order scheme.

            TODO: Update function to get next half timestep if goes into next hour

            Arguments
            ---------
            X, Y: mesh grid.
            x, y: coordinates where to begin the evolution.
            f: the current vector f that will be evolved.
            h: the time step
            dim: 0 for x and 1 for y.

            Returns
            ---------
            interp_value: interpolated value of f(y,x)
            """
            X = self.x_grid
            Y = self.y_grid
            if dim == 0:
                f = self.current_dataset['u_clean'].isel(time=self.time_index).values * 60 * 60
            if dim == 1:
                f = self.current_dataset['v_clean'].isel(time=self.time_index).values * 60 * 60
            try:
                k1 = h * self.bilinear_interpolation(X, Y, f, x, y)
                k2 = h * self.bilinear_interpolation(X, Y, f, x + 0.5 * h, y + 0.5 * k1)
                k3 = h * self.bilinear_interpolation(X, Y, f, x + 0.5 * h, y + 0.5 * k2)
                k4 = h * self.bilinear_interpolation(X, Y, f, x + h, y + k3)
            except ValueError as e:
                print('Error in Interpolation, trying to interpolate a NAN value')
                print(e)
                return

            try:
                if dim == 0:
                    return x + 1. / 6 * k1 + 1. / 3 * k2 + 1. / 3 * k3 + 1. / 6 * k4
                elif dim == 1:
                    return y + 1. / 6 * k1 + 1. / 3 * k2 + 1. / 3 * k3 + 1. / 6 * k4
            except Exception as e:
                print(e)
                sys.exit()

    def bilinear_interpolation(self, X, Y, f, x, y):
        """
        Interpolation methods for estimating surface current values in between grid points. Edge cases are outlined in the
        top of the function and may need to be refactored. NaNs is returned for cases where values can not be interpolated.

        Arguments
        ---------
        X, Y: Coordinate mesh grid
        f: Grid of velocity values that can be accessed as f(j,i) Remeber row, column
        x, y: coordinates to compute interpolation to f(y,x)

        Returns
        ---------
        interp_value: interpolated value of f(y,x)
        """

        # Grid index shape
        M = np.shape(X[:, 0])[0]
        N = np.shape(X[0, :])[0]

        dx, dy = X[0, 1] - X[0, 0], Y[1, 0] - Y[0, 0]
        x_start, y_start = X[0, 0], Y[0, 0]

        # Find the index of each value
        i1, i2 = int((x - x_start) / dx), int((x - x_start) / dx) + 1
        j1, j2 = int((y - y_start) / dy), int((y - y_start) / dy) + 1

        # Boundary Conditions when interpolating near the edge.
        # 1. Eastern boundary
        if (i1 - N) > 1:
            return np.nan
        if i1 >= N - 1 and j1 <= N - 1 and j1 >= 0: # If on the Eastern edge of the boundary
            return f[j1, N - 1]
        if i1 >= N - 1 and j1 <= 0:
            return f[0, N - 1]
        if i1 >= N - 1 and j1 >= N - 1:
            return f[N - 1, N - 1]

        # 2. Western boundary
        if i1 <= 0 and j1 <= N - 1 and j1 >= 0:
            return f[j1, 0]
        if i1 <= 0 and j1 <= 0:
            return f[0, 0]
        if i1 <= 0 and j1 >= N - 1:
            return f[N - 1, 0]

        # 3. Northern boundary
        if j1 >= M - 1 and i1 <= M - 1 and i1 >= 0:
            return f[M - 1, i1]
        if j1 >= N - 1 and i1 <= 0:
            return f[M - 1, 0]

        # 3. Bottom boundary
        if j1 <= 0 and i1 <= M - 1 and i1 >= 0:
            return f[0, i1]
        if j1 <= 0 and i1 >= M - 1:
            return f[M - 1, 0]

        x1, x2 = X[j1, i1], X[j2, i2]
        y1, y2 = Y[j1, i1], Y[j2, i2]

        interp_value = (1 / (x2 - x1) * 1 / (y2 - y1) *
                          (f[j1, i1] * (x2 - x) * (y2 - y) + f[j1, i2] * (x - x1) * (y2 - y)
                           + f[j2, i1] * (x2 - x) * (y - y1) + f[j2, i2] * (x - x1) * (y - y1)))

        return interp_value

    def seed_particles(self, center_coord, radius):
        ''' Create a of cluster of particles within a radius in (km)'''

        x_pos, y_pos = ccrs.UTM(zone=10).transform_point(center_coord[0],center_coord[1], ccrs.PlateCarree())

        col_num = [1,3,5,3,1]
        dx = radius/4 * 1000
        ylevel = (np.arange(max(col_num)) - 3) * dx
        for i, n in enumerate(col_num):
            x = np.arange(n)
            x = x - (n-1)/2
            x = x * dx
            y = np.ones(shape=x.shape) * ylevel[i]
            pos_x = x_pos + x
            pos_y = y_pos + y
            coors = ccrs.PlateCarree().transform_points(ccrs.UTM(zone=10), pos_x, pos_y)
            x_coors = coors[:,0]
            y_coors = coors[:,1]
            for pos in zip(x_coors, y_coors):
                self.add_particle(pos)

class Particle(object):

    def __init__(self,coord):
        self.coordinates = np.array([coord])

    def get_position(self):
        ''' Return latest postion of particle class '''
        return self.coordinates[-1,:]

    def update_position(self, pos):
        ''' Append to coord array '''
        pos = np.array(pos)
        self.coordinates = np.vstack((self.coordinates,pos))

if __name__ == "__main__":
    model = TrackingModel()
    model.draw_map(draw_bathy=True)
    # Debugging Particles
    center_coord = [-122.6, 37.77]
    model.seed_particles(center_coord, radius=3)
    ax = model.geo_axes
#    ax.quiver(model.lon_grid, model.lat_grid, model.current_dataset['u_clean'].isel(time=0), model.current_dataset['v'].isel(time=0))
    ix = np.where(np.isfinite(model.current_dataset['u_clean'].isel(time=0)))
    ixnan = np.where(np.isnan(model.current_dataset['u_clean'].isel(time=0)))
#    ax.scatter(model.lon_grid[ix], model.lat_grid[ix], marker='.', s=10)
    ax.scatter(model.lon_grid[ixnan], model.lat_grid[ixnan], marker='x', s=10, c='r')
#
    for i in range(24*4*3):
        try:
            model.advect_particle()

        except Exception as e:
            print(e)
            print(round(i/4,2),'hours have passed before breaking')
            break
    model.plot_particles(plot_type='animation')


<reponame>sandflow/camdkit
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

# Copyright (c) 2022, Sandflow Consulting LLC
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

'''RED camera reader'''

import csv
import typing
from fractions import Fraction

import camdkit.model
import camdkit.red.cooke as cooke

_LENS_NAME_PIXEL_PITCH_MAP = {
  "RAPTOR 8K VV": 5,
  "MONSTRO 8K VV": 5,
  "KOMODO 6K S35": 4.4,
  "HELIUM 8K S35": 3.65,
  "GEMINI 5K S35": 6,
  "DRAGON": 5
}

def to_clip(meta_3_file: typing.IO, meta_5_file: typing.IO) -> camdkit.model.Clip:
  """Read RED camera metadata into a `Clip`.
  `meta_3_file`: Static camera metadata. CSV file generated using REDline (`REDline --silent --i {camera_file_path} --printMeta 3`)
  `meta_5_file`: Per-frame camera metadata. CSV file generated using REDline (`REDline --silent --i {camera_file_path} --printMeta 5`)
  """

  # read clip metadata
  clip_metadata = next(csv.DictReader(meta_3_file))
  clip = camdkit.model.Clip()

  clip.set_iso(int(clip_metadata['ISO']))

  clip.set_lens_serial_number(clip_metadata["Lens Serial Number"])

  clip.set_white_balance(int(clip_metadata["Kelvin"]))

  clip.set_sensor_pixel_dimensions(
    camdkit.model.SensorPixelDimensions(
      width=int(clip_metadata["Frame Width"]),
      height=int(clip_metadata["Frame Height"])
    )
  )

  pixel_pitch = _LENS_NAME_PIXEL_PITCH_MAP[clip_metadata["Sensor Name"]]
  pix_dims = clip.get_sensor_pixel_dimensions()
  clip.set_sensor_physical_dimensions(
    camdkit.model.SensorPhysicalDimensions(
      width=round(pix_dims.width * pixel_pitch),
      height=round(pix_dims.height * pixel_pitch)
    )
  )

  # read frame metadata
  csv_data = list(csv.DictReader(meta_5_file))

  n_frames = int(clip_metadata["Total Frames"])

  if len(csv_data) != n_frames:
    raise ValueError(f"Inconsistent frame count between header {n_frames} and frame {len(csv_data)} files")

  clip.set_duration(len(csv_data)/Fraction(clip_metadata["FPS"]))

  clip.set_fps(Fraction(clip_metadata["FPS"]))

  clip.set_focal_length(tuple(int(m["Focal Length"]) * 1000 for m in csv_data))

  clip.set_focal_position(tuple(int(m["Focus Distance"]) * 1000 for m in csv_data))

  cooke_metadata = tuple(cooke.from_binary_string(bytes(int(i, 16) for i in m["Cooke Metadata"].split("/"))) for m in csv_data)

  clip.set_entrance_pupil_position(m.entrance_pupil_position for m in cooke_metadata)

  clip.set_iris_position(Fraction(m.aperture_value, 100) for m in cooke_metadata)

  return clip