<reponame>amirdel/dispersion-continua
# Copyright 2017 <NAME>, <EMAIL>
#
# Permission to use, copy, modify, and/or distribute this software for any purpose with or without fee
# is hereby granted, provided that the above copyright notice and this permission notice appear in all
# copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE
# INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE
# FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
# LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION,
# ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import random as random
import bisect as bs
import numpy as np
from py_dp.dispersion.binning import get_cdf
from py_dp.dispersion.second_order_markov import find_1d_bins
from py_dp.dispersion.dispersion_models import dispersionModelGeneral

# These classes are not used now. Just kept for reference.

class dispModelUncorrStencil(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dt, dx_array, x_max,
                 inj_location = "start", verbose = True):
        super(dispModelUncorrStencil,self).__init__(n_particles, n_steps,
                                                    inj_location, verbose)
        self.dx_array = dx_array
        self.dt = dt
        self.table_size = len(dx_array) - 1
        self.x_max = x_max

    def advance_one_step(self, particle_number, current_index):
        x_max = self.x_max
        end_of_domain_reached = False
        dx_array = self.dx_array
        dt = self.dt
        table_size = self.table_size
        rand_ind = random.randint(0,table_size)
        dx = dx_array[rand_ind]
        current_t = self.time_array[particle_number, current_index]
        current_x = self.x_array[particle_number, current_index]
        next_x = current_x + dx
        if next_x > x_max:
            velocity = dx/dt
            distance_to_end = x_max - current_x
            dt = distance_to_end/velocity
            next_x = x_max
            end_of_domain_reached = True
        next_index = current_index + 1
        self.time_array[particle_number,next_index] = current_t + dt
        self.x_array[particle_number,next_index] = next_x
        return end_of_domain_reached

    def follow_specific_paticle(self,particle_number):
        n_steps = self.n_steps
        for step in range(n_steps):
            #x_array, pore_nr_array, time_array entries are changed inside
            #advance_to_next_pore
            current_index = step
            end_flag = self.advance_one_step(particle_number, current_index)
            if end_flag:
                freeze_idx = current_index + 1
                self.freeze_particle(particle_number, freeze_idx)
                break

    def follow_all_particles(self):
        n_particles = self.n_particles
        for particle in range(n_particles):
            self.follow_specific_paticle(particle)

    def freeze_particle(self,particle_number,current_index):
        """
        after a particle gets to the end of the domain it would stay there.
        this function would copy the value at current_idx to all following values for x and time
        """
        self.x_array[particle_number,current_index:] = self.x_array[particle_number,current_index]
        self.time_array[particle_number,current_index:] = self.time_array[particle_number,current_index]
        #self.freeze_time[particle_number] = self.time_array[particle_number,current_index]
        self.last_index_array[particle_number] = current_index

class dispModelCorrelatedStencil(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dt, x_max, trans_matrix,
                 class_velocity, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelCorrelatedStencil,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_velocity = class_velocity
        self.dt = dt
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.cdf_matrix = np.cumsum(trans_matrix, axis=0)

    def draw_from_init_calss_idx(self):
        return bs.bisect(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        cdf = self.cdf_matrix[:, current_class]
        return bs.bisect(cdf, random.random())

    def follow_one_particle(self, particle_number):
        dt = self.dt
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        class_idx = self.draw_from_init_calss_idx()
        next_idx = 0
        v = class_velocity[class_idx]
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x += dt*v
            t += dt
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v = class_velocity[next_idx]
            class_idx = next_idx

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)

class dispModelCorrelatedStencilFix(dispModelCorrelatedStencil):
    def __init__(self, n_particles, n_steps, dt, x_max, trans_matrix,
                 class_velocity, init_class_count, length, inj_location = "start", verbose = True):
        super(dispModelCorrelatedStencilFix,self).__init__(n_particles, n_steps, dt, x_max, trans_matrix,
                 class_velocity, init_class_count, inj_location, verbose)
        self.length = length

    def follow_one_particle(self, particle_number):
        l = self.length
        dt = self.dt
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        # initialize the particle velocity
        class_idx = self.draw_from_init_calss_idx()
        next_idx = 0
        v = class_velocity[class_idx]
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            dx = v*dt
            abs_dx = abs(dx)
            if abs_dx < l:
                length_traveled = 0.0
                while abs(length_traveled) <= l - abs_dx and out_put_idx < idx_max:
                    length_traveled += dx
                    x += dx
                    t += dt
                    x_array[particle_number, out_put_idx] = x
                    t_array[particle_number, out_put_idx] = t
                    out_put_idx += 1
            else:
                x += dt * v
                t += dt
                x_array[particle_number, out_put_idx] = x
                t_array[particle_number, out_put_idx] = t
                out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v = class_velocity[next_idx]
            class_idx = next_idx

class dispModelCorrelatedSpace(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dx, x_max, trans_matrix,
                 class_velocity, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelCorrelatedSpace,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_velocity = class_velocity
        self.dx = dx
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.cdf_matrix = np.cumsum(trans_matrix, axis=0)

    def draw_from_init_calss_idx(self):
        return bs.bisect(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        cdf = self.cdf_matrix[:, current_class]
        return bs.bisect(cdf, random.random())

    def follow_one_particle(self, particle_number):
        dx = self.dx
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        class_idx = self.draw_from_init_calss_idx()
        v = class_velocity[class_idx]
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x += np.sign(v)*dx
            t += dx/abs(v)
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v = class_velocity[next_idx]
            class_idx = next_idx

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)

class dispModelCorrelatedSpaceKang(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dx, x_max, trans_matrix,
                 class_log_edges, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelCorrelatedSpaceKang,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_log_edges = class_log_edges
        self.class_velocity = self.get_class_velocity(class_log_edges)
        self.dx = dx
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.cdf_matrix = np.cumsum(trans_matrix, axis=0)

    def get_class_velocity(self, class_log_edges):
        v_log_edges = self.class_log_edges
        n_class = len(class_log_edges) - 1
        class_velocity = np.zeros(n_class)
        for i in range(n_class):
            log_value = 0.5*(v_log_edges[i] + v_log_edges[i+1])
            class_velocity[i] = np.exp(log_value)
        return class_velocity

    def draw_from_class_velocity(self, idx):
        v_log_edges = self.class_log_edges
        x = random.random()
        log_v = v_log_edges[idx]*x + v_log_edges[idx+1]*(1-x)
        return np.exp(log_v)

    def draw_from_init_calss_idx(self):
        return bs.bisect(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        cdf = self.cdf_matrix[:, current_class]
        return bs.bisect(cdf, random.random())

    def follow_one_particle(self, particle_number):
        dx = self.dx
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        v_class_idx = self.draw_from_init_calss_idx()
        class_idx = 2*v_class_idx
        v = self.draw_from_class_velocity(v_class_idx)
        v_sign = 1.0
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x += v_sign*dx
            t += dx/v
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v_class_idx = np.floor(next_idx/2)
            v_sign = -1.0 + 2.0*((next_idx - 2*v_class_idx) == 0)
            v = self.draw_from_class_velocity(v_class_idx)
            class_idx = next_idx

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)

class dispModelCorrelatedStencilKang(dispersionModelGeneral):
    """
    Class to model plume spreading using a Markov model in time, The velocity is
    binned using the binning strategy in Kang 2010
    """
    def __init__(self, n_particles, n_steps, dt, x_max, trans_matrix,
                 class_log_edges, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelCorrelatedStencilKang,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_log_edges = class_log_edges
        self.dt = dt
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.cdf_matrix = np.cumsum(trans_matrix, axis=0)

    def draw_from_init_calss_idx(self):
        return bs.bisect(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        cdf = self.cdf_matrix[:, current_class]
        return bs.bisect(cdf, random.random())

    def draw_from_class_velocity(self, idx):
        v_log_edges = self.class_log_edges
        x = random.random()
        log_v = v_log_edges[idx]*x + v_log_edges[idx+1]*(1-x)
        return np.exp(log_v)

    def follow_one_particle(self, particle_number):
        dt = self.dt
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        v_class_idx = self.draw_from_init_calss_idx()
        class_idx = 2*v_class_idx
        #v is the abs value of velocity
        v = self.draw_from_class_velocity(v_class_idx)
        v_sign = 1.0
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x += dt*v*v_sign
            t += dt
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v_class_idx = np.floor(next_idx/2)
            v_sign = -1.0 + 2.0*((next_idx - 2*v_class_idx) == 0)
            v = self.draw_from_class_velocity(v_class_idx)
            class_idx = next_idx

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)

class dispModelOrderTwo(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dx, x_max, trans_matrix,
                 class_log_edges, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelOrderTwo,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_log_edges = class_log_edges
        self.class_velocity = self.get_class_velocity(class_log_edges)
        self.dx = dx
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.n_class = np.sqrt(trans_matrix.shape[0])
        self.blocked_particles = []

    def get_class_velocity(self, class_log_edges):
        v_log_edges = self.class_log_edges
        n_class = len(class_log_edges) - 1
        class_velocity = np.zeros(n_class)
        for i in range(n_class):
            log_value = 0.5*(v_log_edges[i] + v_log_edges[i+1])
            class_velocity[i] = np.exp(log_value)
        return class_velocity

    def draw_from_class_velocity(self, idx):
        v_log_edges = self.class_log_edges
        x = random.random()
        log_v = v_log_edges[idx]*x + v_log_edges[idx+1]*(1-x)
        return np.exp(log_v)

    def draw_from_init_calss_idx(self):
        return bs.bisect_right(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        indptr = self.trans_matrix.indptr
        start = indptr[current_class]
        end = indptr[current_class+1]
        rows = self.trans_matrix.indices[start:end]
        values = self.trans_matrix.data[start:end]
        if len(values) == 0:
            return -12
        cdf = get_cdf(values)
        return rows[bs.bisect(cdf, random.random())]

    def advance_x_t(self, v, v_sign, x, t):
        t2 = t + self.dx/v
        x2 = x + v_sign*self.dx
        return x2, t2

    def follow_one_particle(self, particle_number):
        n_class = self.n_class
        dx = self.dx
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        #class_idx is the index of the 2d class
        class_idx = self.draw_from_init_calss_idx()
        #i, ip1 are indices of (abs(v), sgn(v)) classes
        i, ip1 = find_1d_bins(class_idx, n_class)
        v_class_idx = np.floor(i/2)
        v_sign = -1.0 + 2.0*((i - 2*v_class_idx) == 0)
        v = self.draw_from_class_velocity(v_class_idx)
        x,t = self.advance_x_t(v, v_sign, x, t)
        x_array[particle_number, out_put_idx] = x
        t_array[particle_number, out_put_idx] = t
        out_put_idx += 1
        v_class_idx = np.floor(ip1/2)
        v_sign = -1.0 + 2.0*((ip1 - 2*v_class_idx) == 0)
        v = self.draw_from_class_velocity(v_class_idx)
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x,t = self.advance_x_t(v, v_sign, x,t)
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            if next_idx == -12:
                self.blocked_particles.append(particle_number)
                return
            t1, t2 = find_1d_bins(next_idx, n_class)
            class_idx = next_idx
            i, ip1 = find_1d_bins(class_idx, n_class)
            v_class_idx = np.floor(ip1/2)
            v_sign = -1.0 + 2.0*((ip1 - 2*v_class_idx) == 0)
            v = self.draw_from_class_velocity(v_class_idx)

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)
        print "removing blocked particles: ", len(self.blocked_particles)
        idx_array = np.array(range(self.n_particles))
        blocked = np.array(self.blocked_particles)
        idx_diff = np.setdiff1d(idx_array, blocked)
        self.x_array = self.x_array[idx_diff]
        self.time_array = self.time_array[idx_diff]
        self.n_particles -= len(self.blocked_particles)


class dispModelTime3d(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dt, x_max, trans_matrix,
                 mapping, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelTime3d,self).__init__(n_particles, n_steps, inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.mapping = mapping
        self.dt = dt
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.n_class = np.sqrt(trans_matrix.shape[0])
        self.blocked_particles = []

    def draw_from_class_velocity(self, idx):
        v_log_edges = self.class_log_edges
        x = random.random()
        log_v = v_log_edges[idx]*x + v_log_edges[idx+1]*(1-x)
        return np.exp(log_v)

    def draw_from_init_calss_idx(self):
        return bs.bisect_right(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        indptr = self.trans_matrix.indptr
        start = indptr[current_class]
        end = indptr[current_class+1]
        rows = self.trans_matrix.indices[start:end]
        values = self.trans_matrix.data[start:end]
        if len(values) == 0:
            return -12
        cdf = get_cdf(values)
        return rows[bs.bisect_left(cdf, random.random())]

    def advance_x_t(self, v, v_sign, freq, x, t):
        dt = self.dt
        dx = v_sign*v*dt
        if freq>1:
            t2 = np.arange(t, t + freq*dt, dt)
            x2 = x + np.arange(1,1+freq)*dx
        else:
            t2 = t + dt
            x2 = x + v_sign*v*dt
        return x2, t2

    def follow_one_particle(self, particle_number):
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        #class_idx is the index of the 2d class
        class_idx = self.draw_from_init_calss_idx()
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            v_1d_class, v_sign, freq = self.mapping.find_absvclass_sgn_freq(class_idx)
            v = self.mapping.draw_from_class_velocity(v_1d_class)
            x_new, t_new = self.advance_x_t(v, v_sign, freq, x, t)
            if freq>1:
                end_idx = min(out_put_idx+freq, idx_max)
                len_idx = end_idx - out_put_idx
                x_array[particle_number, out_put_idx:end_idx] = x_new[:len_idx]
                t_array[particle_number, out_put_idx:end_idx] = t_new[:len_idx]
                out_put_idx += freq
                x = x_new[-1]
                t = t_new[-1]
            else:
                x_array[particle_number, out_put_idx] = x_new
                t_array[particle_number, out_put_idx] = t_new
                out_put_idx += 1
                x = x_new
                t = t_new
            next_idx = self.choose_next_class(class_idx)
            class_idx = next_idx
            if next_idx == -12:
                self.blocked_particles.append(particle_number)
                return

    def follow_all_particles(self):
        for i in range(self.n_particles):
            if not np.mod(i,200):
                print 'particle number: ', i
            self.follow_one_particle(i)
        print "removing blocked particles: ", len(self.blocked_particles)
        idx_array = np.array(range(self.n_particles))
        blocked = np.array(self.blocked_particles)
        idx_diff = np.setdiff1d(idx_array, blocked)
        self.x_array = self.x_array[idx_diff]
        self.time_array = self.time_array[idx_diff]
        self.n_particles -= len(self.blocked_particles)

# Copyright 2017 <NAME>, <EMAIL>
#
# Permission to use, copy, modify, and/or distribute this software for any purpose with or without fee
# is hereby granted, provided that the above copyright notice and this permission notice appear in all
# copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE
# INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE
# FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
# LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION,
# ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import random as random
import bisect as bs
import numpy as np
from py_dp.dispersion.binning import get_cdf
from py_dp.dispersion.second_order_markov import find_1d_bins
from py_dp.dispersion.dispersion_models import dispersionModelGeneral

# These classes are not used now. Just kept for reference.

class dispModelUncorrStencil(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dt, dx_array, x_max,
                 inj_location = "start", verbose = True):
        super(dispModelUncorrStencil,self).__init__(n_particles, n_steps,
                                                    inj_location, verbose)
        self.dx_array = dx_array
        self.dt = dt
        self.table_size = len(dx_array) - 1
        self.x_max = x_max

    def advance_one_step(self, particle_number, current_index):
        x_max = self.x_max
        end_of_domain_reached = False
        dx_array = self.dx_array
        dt = self.dt
        table_size = self.table_size
        rand_ind = random.randint(0,table_size)
        dx = dx_array[rand_ind]
        current_t = self.time_array[particle_number, current_index]
        current_x = self.x_array[particle_number, current_index]
        next_x = current_x + dx
        if next_x > x_max:
            velocity = dx/dt
            distance_to_end = x_max - current_x
            dt = distance_to_end/velocity
            next_x = x_max
            end_of_domain_reached = True
        next_index = current_index + 1
        self.time_array[particle_number,next_index] = current_t + dt
        self.x_array[particle_number,next_index] = next_x
        return end_of_domain_reached

    def follow_specific_paticle(self,particle_number):
        n_steps = self.n_steps
        for step in range(n_steps):
            #x_array, pore_nr_array, time_array entries are changed inside
            #advance_to_next_pore
            current_index = step
            end_flag = self.advance_one_step(particle_number, current_index)
            if end_flag:
                freeze_idx = current_index + 1
                self.freeze_particle(particle_number, freeze_idx)
                break

    def follow_all_particles(self):
        n_particles = self.n_particles
        for particle in range(n_particles):
            self.follow_specific_paticle(particle)

    def freeze_particle(self,particle_number,current_index):
        """
        after a particle gets to the end of the domain it would stay there.
        this function would copy the value at current_idx to all following values for x and time
        """
        self.x_array[particle_number,current_index:] = self.x_array[particle_number,current_index]
        self.time_array[particle_number,current_index:] = self.time_array[particle_number,current_index]
        #self.freeze_time[particle_number] = self.time_array[particle_number,current_index]
        self.last_index_array[particle_number] = current_index

class dispModelCorrelatedStencil(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dt, x_max, trans_matrix,
                 class_velocity, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelCorrelatedStencil,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_velocity = class_velocity
        self.dt = dt
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.cdf_matrix = np.cumsum(trans_matrix, axis=0)

    def draw_from_init_calss_idx(self):
        return bs.bisect(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        cdf = self.cdf_matrix[:, current_class]
        return bs.bisect(cdf, random.random())

    def follow_one_particle(self, particle_number):
        dt = self.dt
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        class_idx = self.draw_from_init_calss_idx()
        next_idx = 0
        v = class_velocity[class_idx]
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x += dt*v
            t += dt
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v = class_velocity[next_idx]
            class_idx = next_idx

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)

class dispModelCorrelatedStencilFix(dispModelCorrelatedStencil):
    def __init__(self, n_particles, n_steps, dt, x_max, trans_matrix,
                 class_velocity, init_class_count, length, inj_location = "start", verbose = True):
        super(dispModelCorrelatedStencilFix,self).__init__(n_particles, n_steps, dt, x_max, trans_matrix,
                 class_velocity, init_class_count, inj_location, verbose)
        self.length = length

    def follow_one_particle(self, particle_number):
        l = self.length
        dt = self.dt
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        # initialize the particle velocity
        class_idx = self.draw_from_init_calss_idx()
        next_idx = 0
        v = class_velocity[class_idx]
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            dx = v*dt
            abs_dx = abs(dx)
            if abs_dx < l:
                length_traveled = 0.0
                while abs(length_traveled) <= l - abs_dx and out_put_idx < idx_max:
                    length_traveled += dx
                    x += dx
                    t += dt
                    x_array[particle_number, out_put_idx] = x
                    t_array[particle_number, out_put_idx] = t
                    out_put_idx += 1
            else:
                x += dt * v
                t += dt
                x_array[particle_number, out_put_idx] = x
                t_array[particle_number, out_put_idx] = t
                out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v = class_velocity[next_idx]
            class_idx = next_idx

class dispModelCorrelatedSpace(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dx, x_max, trans_matrix,
                 class_velocity, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelCorrelatedSpace,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_velocity = class_velocity
        self.dx = dx
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.cdf_matrix = np.cumsum(trans_matrix, axis=0)

    def draw_from_init_calss_idx(self):
        return bs.bisect(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        cdf = self.cdf_matrix[:, current_class]
        return bs.bisect(cdf, random.random())

    def follow_one_particle(self, particle_number):
        dx = self.dx
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        class_idx = self.draw_from_init_calss_idx()
        v = class_velocity[class_idx]
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x += np.sign(v)*dx
            t += dx/abs(v)
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v = class_velocity[next_idx]
            class_idx = next_idx

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)

class dispModelCorrelatedSpaceKang(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dx, x_max, trans_matrix,
                 class_log_edges, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelCorrelatedSpaceKang,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_log_edges = class_log_edges
        self.class_velocity = self.get_class_velocity(class_log_edges)
        self.dx = dx
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.cdf_matrix = np.cumsum(trans_matrix, axis=0)

    def get_class_velocity(self, class_log_edges):
        v_log_edges = self.class_log_edges
        n_class = len(class_log_edges) - 1
        class_velocity = np.zeros(n_class)
        for i in range(n_class):
            log_value = 0.5*(v_log_edges[i] + v_log_edges[i+1])
            class_velocity[i] = np.exp(log_value)
        return class_velocity

    def draw_from_class_velocity(self, idx):
        v_log_edges = self.class_log_edges
        x = random.random()
        log_v = v_log_edges[idx]*x + v_log_edges[idx+1]*(1-x)
        return np.exp(log_v)

    def draw_from_init_calss_idx(self):
        return bs.bisect(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        cdf = self.cdf_matrix[:, current_class]
        return bs.bisect(cdf, random.random())

    def follow_one_particle(self, particle_number):
        dx = self.dx
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        v_class_idx = self.draw_from_init_calss_idx()
        class_idx = 2*v_class_idx
        v = self.draw_from_class_velocity(v_class_idx)
        v_sign = 1.0
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x += v_sign*dx
            t += dx/v
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v_class_idx = np.floor(next_idx/2)
            v_sign = -1.0 + 2.0*((next_idx - 2*v_class_idx) == 0)
            v = self.draw_from_class_velocity(v_class_idx)
            class_idx = next_idx

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)

class dispModelCorrelatedStencilKang(dispersionModelGeneral):
    """
    Class to model plume spreading using a Markov model in time, The velocity is
    binned using the binning strategy in Kang 2010
    """
    def __init__(self, n_particles, n_steps, dt, x_max, trans_matrix,
                 class_log_edges, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelCorrelatedStencilKang,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_log_edges = class_log_edges
        self.dt = dt
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.cdf_matrix = np.cumsum(trans_matrix, axis=0)

    def draw_from_init_calss_idx(self):
        return bs.bisect(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        cdf = self.cdf_matrix[:, current_class]
        return bs.bisect(cdf, random.random())

    def draw_from_class_velocity(self, idx):
        v_log_edges = self.class_log_edges
        x = random.random()
        log_v = v_log_edges[idx]*x + v_log_edges[idx+1]*(1-x)
        return np.exp(log_v)

    def follow_one_particle(self, particle_number):
        dt = self.dt
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        v_class_idx = self.draw_from_init_calss_idx()
        class_idx = 2*v_class_idx
        #v is the abs value of velocity
        v = self.draw_from_class_velocity(v_class_idx)
        v_sign = 1.0
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x += dt*v*v_sign
            t += dt
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            v_class_idx = np.floor(next_idx/2)
            v_sign = -1.0 + 2.0*((next_idx - 2*v_class_idx) == 0)
            v = self.draw_from_class_velocity(v_class_idx)
            class_idx = next_idx

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)

class dispModelOrderTwo(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dx, x_max, trans_matrix,
                 class_log_edges, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelOrderTwo,self).__init__(n_particles, n_steps,
                                                        inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.class_log_edges = class_log_edges
        self.class_velocity = self.get_class_velocity(class_log_edges)
        self.dx = dx
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.n_class = np.sqrt(trans_matrix.shape[0])
        self.blocked_particles = []

    def get_class_velocity(self, class_log_edges):
        v_log_edges = self.class_log_edges
        n_class = len(class_log_edges) - 1
        class_velocity = np.zeros(n_class)
        for i in range(n_class):
            log_value = 0.5*(v_log_edges[i] + v_log_edges[i+1])
            class_velocity[i] = np.exp(log_value)
        return class_velocity

    def draw_from_class_velocity(self, idx):
        v_log_edges = self.class_log_edges
        x = random.random()
        log_v = v_log_edges[idx]*x + v_log_edges[idx+1]*(1-x)
        return np.exp(log_v)

    def draw_from_init_calss_idx(self):
        return bs.bisect_right(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        indptr = self.trans_matrix.indptr
        start = indptr[current_class]
        end = indptr[current_class+1]
        rows = self.trans_matrix.indices[start:end]
        values = self.trans_matrix.data[start:end]
        if len(values) == 0:
            return -12
        cdf = get_cdf(values)
        return rows[bs.bisect(cdf, random.random())]

    def advance_x_t(self, v, v_sign, x, t):
        t2 = t + self.dx/v
        x2 = x + v_sign*self.dx
        return x2, t2

    def follow_one_particle(self, particle_number):
        n_class = self.n_class
        dx = self.dx
        class_velocity = self.class_velocity
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        #class_idx is the index of the 2d class
        class_idx = self.draw_from_init_calss_idx()
        #i, ip1 are indices of (abs(v), sgn(v)) classes
        i, ip1 = find_1d_bins(class_idx, n_class)
        v_class_idx = np.floor(i/2)
        v_sign = -1.0 + 2.0*((i - 2*v_class_idx) == 0)
        v = self.draw_from_class_velocity(v_class_idx)
        x,t = self.advance_x_t(v, v_sign, x, t)
        x_array[particle_number, out_put_idx] = x
        t_array[particle_number, out_put_idx] = t
        out_put_idx += 1
        v_class_idx = np.floor(ip1/2)
        v_sign = -1.0 + 2.0*((ip1 - 2*v_class_idx) == 0)
        v = self.draw_from_class_velocity(v_class_idx)
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            x,t = self.advance_x_t(v, v_sign, x,t)
            x_array[particle_number, out_put_idx] = x
            t_array[particle_number, out_put_idx] = t
            out_put_idx += 1
            next_idx = self.choose_next_class(class_idx)
            if next_idx == -12:
                self.blocked_particles.append(particle_number)
                return
            t1, t2 = find_1d_bins(next_idx, n_class)
            class_idx = next_idx
            i, ip1 = find_1d_bins(class_idx, n_class)
            v_class_idx = np.floor(ip1/2)
            v_sign = -1.0 + 2.0*((ip1 - 2*v_class_idx) == 0)
            v = self.draw_from_class_velocity(v_class_idx)

    def follow_all_particles(self):
        for i in range(self.n_particles):
            self.follow_one_particle(i)
        print "removing blocked particles: ", len(self.blocked_particles)
        idx_array = np.array(range(self.n_particles))
        blocked = np.array(self.blocked_particles)
        idx_diff = np.setdiff1d(idx_array, blocked)
        self.x_array = self.x_array[idx_diff]
        self.time_array = self.time_array[idx_diff]
        self.n_particles -= len(self.blocked_particles)


class dispModelTime3d(dispersionModelGeneral):
    def __init__(self, n_particles, n_steps, dt, x_max, trans_matrix,
                 mapping, init_class_count, inj_location = "start",
                 verbose = True):
        super(dispModelTime3d,self).__init__(n_particles, n_steps, inj_location, verbose)
        self.trans_matrix = trans_matrix
        self.init_class_count = init_class_count
        self.mapping = mapping
        self.dt = dt
        self.x_max = x_max
        self.init_class_cdf = get_cdf(init_class_count)
        self.n_class = np.sqrt(trans_matrix.shape[0])
        self.blocked_particles = []

    def draw_from_class_velocity(self, idx):
        v_log_edges = self.class_log_edges
        x = random.random()
        log_v = v_log_edges[idx]*x + v_log_edges[idx+1]*(1-x)
        return np.exp(log_v)

    def draw_from_init_calss_idx(self):
        return bs.bisect_right(self.init_class_cdf, random.random())

    def choose_next_class(self, current_class):
        indptr = self.trans_matrix.indptr
        start = indptr[current_class]
        end = indptr[current_class+1]
        rows = self.trans_matrix.indices[start:end]
        values = self.trans_matrix.data[start:end]
        if len(values) == 0:
            return -12
        cdf = get_cdf(values)
        return rows[bs.bisect_left(cdf, random.random())]

    def advance_x_t(self, v, v_sign, freq, x, t):
        dt = self.dt
        dx = v_sign*v*dt
        if freq>1:
            t2 = np.arange(t, t + freq*dt, dt)
            x2 = x + np.arange(1,1+freq)*dx
        else:
            t2 = t + dt
            x2 = x + v_sign*v*dt
        return x2, t2

    def follow_one_particle(self, particle_number):
        x_array = self.x_array
        t_array = self.time_array
        x = 0.0
        t = 0.0
        out_put_idx = 1
        #initialize the particle velocity
        #class_idx is the index of the 2d class
        class_idx = self.draw_from_init_calss_idx()
        idx_max = self.n_steps + 1
        while out_put_idx < idx_max:
            v_1d_class, v_sign, freq = self.mapping.find_absvclass_sgn_freq(class_idx)
            v = self.mapping.draw_from_class_velocity(v_1d_class)
            x_new, t_new = self.advance_x_t(v, v_sign, freq, x, t)
            if freq>1:
                end_idx = min(out_put_idx+freq, idx_max)
                len_idx = end_idx - out_put_idx
                x_array[particle_number, out_put_idx:end_idx] = x_new[:len_idx]
                t_array[particle_number, out_put_idx:end_idx] = t_new[:len_idx]
                out_put_idx += freq
                x = x_new[-1]
                t = t_new[-1]
            else:
                x_array[particle_number, out_put_idx] = x_new
                t_array[particle_number, out_put_idx] = t_new
                out_put_idx += 1
                x = x_new
                t = t_new
            next_idx = self.choose_next_class(class_idx)
            class_idx = next_idx
            if next_idx == -12:
                self.blocked_particles.append(particle_number)
                return

    def follow_all_particles(self):
        for i in range(self.n_particles):
            if not np.mod(i,200):
                print 'particle number: ', i
            self.follow_one_particle(i)
        print "removing blocked particles: ", len(self.blocked_particles)
        idx_array = np.array(range(self.n_particles))
        blocked = np.array(self.blocked_particles)
        idx_diff = np.setdiff1d(idx_array, blocked)
        self.x_array = self.x_array[idx_diff]
        self.time_array = self.time_array[idx_diff]
        self.n_particles -= len(self.blocked_particles)
# Copyright 2017 <NAME>, <EMAIL> # # Permission to use, copy, modify, and/or distribute this software for any purpose with or without fee # is hereby granted, provided that the above copyright notice and this permission notice appear in all # copies. # # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE # INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE # FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM # LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, # ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. # These classes are not used now. Just kept for reference. #x_array, pore_nr_array, time_array entries are changed inside #advance_to_next_pore after a particle gets to the end of the domain it would stay there.         this function would copy the value at current_idx to all following values for x and time #self.freeze_time[particle_number] = self.time_array[particle_number,current_index] #initialize the particle velocity # initialize the particle velocity #initialize the particle velocity #initialize the particle velocity Class to model plume spreading using a Markov model in time, The velocity is     binned using the binning strategy in Kang 2010 #initialize the particle velocity #v is the abs value of velocity #initialize the particle velocity #class_idx is the index of the 2d class #i, ip1 are indices of (abs(v), sgn(v)) classes #initialize the particle velocity #class_idx is the index of the 2d class