text:
# -*- coding: utf-8 -*-
"""
Created on Mon Apr 11 13:22:59 2016

Cellular Automata: Regolith Formation
Using CellLab

@author: Katherine
"""

import time
import matplotlib
import numpy as np
from landlab import RasterModelGrid
from landlab.ca.celllab_cts import Transition, CAPlotter
from landlab.ca.raster_cts import RasterCTS
from fracture_grid import make_frac_grid

def setup_transition_list():
    """
    Creates and returns a list of Transition() objects to represent state
    transitions for an unbiased random walk.

    Parameters
    ----------
    (none)

    Returns
    -------
    xn_list : list of Transition objects
            List of objects that encode information about the link-state transitions.

    Notes
    -----
    State 0 represents fluid and state 1 represents a particle (such as a
    sediment grain, tea leaf, or solute molecule).

    The states and transitions are as follows:

    Pair state      Transition to       Process             Rate (cells/s)
    ==========      =============       =======             ==============
    0 (0-0)         (none)              -                   -
    1 (1-0)         2 (1-1)             left/down motion    10.0
    2 (0-1)         1 (1-1)             right/up motion     10.0
    3 (1-1)         (none)              -                   -

    """

    # Create an empty transition list
    xn_list = []

    # Append two transitions to the list.
    # Note that the arguments to the Transition() object constructor are:
    #  - Tuple representing starting pair state
    #    (left/bottom cell, right/top cell, orientation)
    #  - Tuple representing new pair state
    #    (left/bottom cell, right/top cell, orientation)
    #  - Transition rate (cells per time step, in this case 1 sec)
    #  - Name for transition
    xn_list.append( Transition((1,0,0), (1,1,0), 1., 'saprolite formation') )
    xn_list.append( Transition((0,1,0), (1,1,0), 1., 'saprolite formation') )
    return xn_list

# INITIALIZE
# User-defined parameters
nr = 300  # number of rows in grid
nc = 300  # number of columns in grid
plot_interval = 0.5   # time interval for plotting, sec
run_duration = 20.0   # duration of run, sec
report_interval = 10.0  # report interval, in real-time seconds

# Remember the clock time, and calculate when we next want to report
# progress.
current_real_time = time.time()
next_report = current_real_time + report_interval

# Create grid
mg = RasterModelGrid(nr, nc, 1.0)

# Make the boundaries be walls
mg.set_closed_boundaries_at_grid_edges(True, True, True, True)

#Create a node-state dictionary
ns_dict = { 0 : 'rock', 1 : 'saprolite' }

#Create the transition list
xn_list = setup_transition_list()

# Create the node-state array and attach it to the grid
node_state_grid = mg.add_zeros('node', 'node_state_map', dtype=int)

# Initialize the node-state array: here, the initial condition is a set of fracture lines.
node_state_grid.shape=((mg.number_of_node_rows,mg.number_of_node_columns))
node_state_grid=make_frac_grid(10, numrows=nr, numcols=nc, model_grid=node_state_grid)

# For visual display purposes, set all boundary nodes to fluid
node_state_grid[mg.closed_boundary_nodes] = 0

# Create the CA model
ca = RasterCTS(mg, ns_dict, xn_list, node_state_grid)

# Set up colors for plotting
grain = '#5F594D'
fluid = '#D0E4F2'
clist = [fluid,grain]
my_cmap = matplotlib.colors.ListedColormap(clist)

# Create a CAPlotter object for handling screen display
ca_plotter = CAPlotter(ca, cmap=my_cmap)

# Plot the initial grid
ca_plotter.update_plot()

# RUN
current_time = 0.0
while current_time < run_duration:

# Once in a while, print out simulation real time to let the user
# know that the sim is running ok
    current_real_time = time.time()
    if current_real_time >= next_report:
        print('Current simulation time '+str(current_time)+'  \
                   ('+str(int(100*current_time/run_duration))+'%)')
        next_report = current_real_time + report_interval

# Run the model forward in time until the next output step
    ca.run(current_time+plot_interval, ca.node_state, plot_each_transition=False)
    current_time += plot_interval

# Plot the current grid
    ca_plotter.update_plot()

ca_plotter.finalize()


text:
"""
Two Point Functions Framework.

"""

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

from yt.utilities.on_demand_imports import _h5py as h5py
import numpy as np

from yt.funcs import mylog
from yt.utilities.performance_counters import yt_counters
from yt.utilities.parallel_tools.parallel_analysis_interface import ParallelAnalysisInterface, parallel_blocking_call, parallel_root_only

try:
    from yt.utilities.kdtree.api import \
        fKD, free_tree, create_tree
except ImportError:
    mylog.debug("The Fortran kD-Tree did not import correctly.")

import math
import inspect
import time
from collections import defaultdict

sep = 12

class TwoPointFunctions(ParallelAnalysisInterface):
    r""" Initialize a two point functions object.

    Parameters
    ----------
    total_values : Integer
        How many total (global) pair calculations to run for each of the
        functions specified. Default: 1000000.
    comm_size : Integer
        How entries are sent during communication. Default: 10000.
    length_type : String
        Controls the even spacing of the rulers lengths in
        logarithmic or linear space, set by "log" or "lin", respectively.
        Default: "lin".
    length_number : Integer
        Sets how many lengths to create, evenly spaced by the above
        parameter. Default: 10.
    length_range : Float
        A min/max pair for the range of values to search the over
        the simulated volume. Default: [sqrt(3)dx, 1/2*shortest box edge],
        where dx is the smallest grid cell size.
    vol_ratio : Integer
        How to multiply-assign subvolumes to the parallel
        tasks. This number must be an integer factor of the total number of tasks or
        very bad things will happen. The default value of 1 will assign one task
        to each subvolume, and there will be an equal number of subvolumes as tasks.
        A value of 2 will assign two tasks to each subvolume and there will be
        one-half as many subvolumes as tasks.
        A value equal to the number of parallel tasks will result in each task
        owning a complete copy of all the fields data, meaning each task will be
        operating on the identical full volume.
        Setting it to -1 will automatically adjust it such that each task
        owns the entire volume. Default = 1.
    salt : Integer
        A number that will be added to the random number generator
        seed. Use this if a different random series of numbers is desired when
        keeping everything else constant from this set: (MPI task count,
        number of ruler lengths, ruler min/max, number of functions,
        number of point pairs per ruler length). Default = 0.
    theta : Float
        For random pairs of points, the second point is found by traversing
        a distance along a ray set by the angle (phi, theta) from the first
        point. To keep this angle constant, set ``theta`` to a value in the
        range [0, pi]. Default = None, which will randomize theta for
        every pair of points.
    phi : Float
        Similar to theta above, but the range of values is [0, 2*pi).
        Default = None, which will randomize phi for every pair of points.

    Examples
    --------
    >>> tpf = TwoPointFunctions(ds, ["velocity_x", "velocity_y", "velocity_z"],
    ... total_values=1e5, comm_size=10000,
    ... length_number=10, length_range=[1./128, .5],
    ... length_type="log")
    """
    def __init__(self, ds, fields, left_edge=None, right_edge=None,
            total_values=1000000, comm_size=10000, length_type="lin",
            length_number=10, length_range=None, vol_ratio = 1,
            salt=0, theta=None, phi=None):
        ParallelAnalysisInterface.__init__(self)
        try:
            fKD
        except NameError:
            raise ImportError("You need to install the Forthon kD-Tree")
        self._fsets = []
        self.fields = fields
        self.constant_theta = theta
        self.constant_phi = phi
        # MPI stuff.
        self.size = self.comm.size
        self.mine = self.comm.rank
        self.vol_ratio = vol_ratio
        if self.vol_ratio == -1:
            self.vol_ratio = self.size
        self.total_values = int(total_values / self.size)
        # For communication.
        self.recv_hooks = []
        self.send_hooks = []
        self.done_hooks = []
        self.comm_size = min(int(comm_size), self.total_values)
        self.ds = ds
        self.nlevels = ds.index.max_level
        self.period = self.ds.domain_right_edge - self.ds.domain_left_edge
        self.min_edge = min(self.period)
        self.index = ds.index
        self.center = (ds.domain_right_edge + ds.domain_left_edge)/2.0
        # Figure out the range of ruler lengths.
        if length_range is None:
            length_range = [math.sqrt(3) * self.ds.index.get_smallest_dx(),
                self.min_edge/2.]
        else:
            if len(length_range) != 2:
                raise ValueError("length_range must have two values.")
            if length_range[1] <= length_range[0]:
                raise ValueError("length_range[1] must be larger than length_range[0]")
            if length_range[1] > self.min_edge/2.:
                length_range[1] = self.min_edge/2.
                mylog.info("Automatically adjusting length_range[1] to half the shortest box edge.")
        if length_range[0] == -1 or length_range[0] == -1.:
            mylog.info("Automatically adjusting length_range[0] to %1.5e." % \
                (math.sqrt(3) * self.ds.index.get_smallest_dx()))
            length_range[0] = math.sqrt(3) * self.ds.index.get_smallest_dx()
        # Make the list of ruler lengths.
        if length_type == "lin":
            self.lengths = np.linspace(length_range[0], length_range[1],
                length_number)
        elif length_type == "log":
            self.lengths = np.logspace(math.log10(length_range[0]),
                math.log10(length_range[1]), length_number)
        else:
            # Something went wrong.
            raise SyntaxError("length_type is either \"lin\" or \"log\".")
        # Subdivide the volume.
        if not left_edge or not right_edge:
            self.left_edge = self.ds.domain_left_edge
            self.right_edge = self.ds.domain_right_edge
            # This ds business below has to do with changes made for halo
            # finding on subvolumes and serves no purpose here except
            # compatibility. This is not the best policy, if I'm honest.
            ds = ds.region([0.]*3, self.left_edge, self.right_edge)
            padded, self.LE, self.RE, self.ds = \
            self.partition_index_3d(ds = ds, padding=0.,
                rank_ratio = self.vol_ratio)
        else:
            self.left_edge = left_edge
            self.right_edge = right_edge
            # We do this twice, first with no 'buffer' to get the unbuffered
            # self.LE/RE, and then second to get a buffered self.ds.
            padded, self.LE, self.RE, temp = \
                self.partition_region_3d(left_edge, right_edge,
                    rank_ratio=self.vol_ratio)
            padded, temp, temp, self.ds = \
                self.partition_region_3d(left_edge - self.lengths[-1], \
                right_edge + self.lengths[-1], rank_ratio=self.vol_ratio)
        mylog.info("LE %s RE %s %s" % (str(self.LE), str(self.RE), str(self.ds)))
        self.width = self.ds.right_edge - self.ds.left_edge
        self.mt = np.random.mtrand.RandomState(seed = 1234 * self.mine + salt)

    def add_function(self, function, out_labels, sqrt, corr_norm=None):
        r"""Add a function to the list that will be evaluated at the
        generated pairs of points.

        Parameters
        ----------
        function : Function
            The two point function of the form fcn(a, b, r1, r2, vec).
        out_labels : List of strings
            A list of strings labeling the outputs of the function.
        sqrt : List of booleans
            A list of booleans which when True will square-root the corresponding
            element of the output in the text output (write_out_means()).
        corr_norm : Float
            Used when calculating two point correlations. If set, the output
            of the function is divided by this number. Default = None.

        Examples
        --------
        >>> f1 = tpf.add_function(function=rms_vel, out_labels=['RMSvdiff'],
        ... sqrt=[True])
        """
        fargs = inspect.getargspec(function)
        if len(fargs.args) != 5:
            raise SyntaxError("The function %s needs five arguments." %\
                function.__name__)
        out_labels = list(out_labels)
        if len(out_labels) < 1:
            raise SyntaxError("Please specify at least one out_labels for function %s." %\
                function.__name__)
        sqrt = list(sqrt)
        if len(sqrt) != len(out_labels):
            raise SyntaxError("Please have the same number of elements in out_labels as in sqrt for function %s." %\
                function.__name__)
        self._fsets.append(FcnSet(self, function, self.min_edge,
            out_labels, sqrt,corr_norm))
        return self._fsets[-1]

    def __getitem__(self, key):
        return self._fsets[key]

    def run_generator(self):
        r"""After all the functions have been added, run the generator.

        Examples
        --------
        >>> tpf.run_generator()
        """
        yt_counters("run_generator")
        # We need a function!
        if len(self._fsets) == 0:
            mylog.error("You need to add at least one function!")
            return None
        # Do all the startup tasks to get the grid points.
        if self.nlevels == 0:
            yt_counters("build_sort")
            self._build_sort_array()
            self.sort_done = False
            yt_counters("build_sort")
        else:
            yt_counters("init_kd_tree")
            self._init_kd_tree()
            self.sort_done = True
            yt_counters("init_kd_tree")
        # Store the fields.
        self.stored_fields = {}
        yt_counters("getting data")
        for field in self.fields:
            self.stored_fields[field] = self.ds[field].copy()
        self.ds.clear_data()
        # If the arrays haven't been sorted yet and need to be, do that.
        if not self.sort_done:
            for field in self.fields:
                self.stored_fields[field] = self.stored_fields[field][self.sort]
            del self.sort
            self.sort_done = True
        yt_counters("getting data")
        self._build_fields_vals()
        yt_counters("big loop over lengths")
        t_waiting = 0.
        for bigloop, length in enumerate(self.lengths):
            self._build_points_array()
            if self.mine == 0:
                mylog.info("Doing length %1.5e" % length)
            # Things stop when this value below equals total_values.
            self.generated_points = 0
            self.gen_array = np.zeros(self.size, dtype='int64')
            self.comm_cycle_count = 0
            self.final_comm_cycle_count = 0
            self.sent_done = False
            self._setup_done_hooks_on_root()
            # While everyone else isn't done or I'm not done, we loop.
            while self._should_cycle():
                self._setup_recv_arrays()
                self._send_arrays()
                t0 = time.time()
                self.comm.mpi_Request_Waitall(self.send_hooks)
                self.comm.mpi_Request_Waitall(self.recv_hooks)
                t1 = time.time()
                t_waiting += (t1-t0)
                if (self.recv_points < -1.).any() or (self.recv_points > 1.).any(): # or \
                    #(np.abs(np.log10(np.abs(self.recv_points))) > 20).any():
                    raise ValueError("self.recv_points is no good!")
                self.points = self.recv_points.copy()
                self.fields_vals = self.recv_fields_vals.copy()
                self.gen_array = self.recv_gen_array.copy()
                self._eval_points(length)
                self.gen_array[self.mine] = self.generated_points
                self.comm_cycle_count += 1
                if self.generated_points == self.total_values:
                    self._send_done_to_root()
            if self.mine == 0:
                mylog.info("Length (%d of %d) %1.5e took %d communication cycles to complete." % \
                (bigloop+1, len(self.lengths), length, self.comm_cycle_count))
        yt_counters("big loop over lengths")
        if self.nlevels >= 1:
            del fKD.pos, fKD.qv_many, fKD.nn_tags
            free_tree(0) # Frees the kdtree object.
        yt_counters("allsum")
        self._allsum_bin_hits()
        mylog.info("Spent %f seconds waiting for communication." % t_waiting)
        yt_counters("allsum")
        yt_counters("run_generator")

    def _init_kd_tree(self):
        """
        Builds the kd tree of grid center points.
        """
        # Grid cell centers.
        mylog.info("Multigrid: Building kD-Tree.")
        xp = self.ds["x"]
        yp = self.ds["y"]
        zp = self.ds["z"]
        fKD.pos = np.asfortranarray(np.empty((3,xp.size), dtype='float64'))
        # Normalize the grid points only within the kdtree.
        fKD.pos[0, :] = xp[:] / self.period[0]
        fKD.pos[1, :] = yp[:] / self.period[1]
        fKD.pos[2, :] = zp[:] / self.period[2]
        fKD.nn = 1
        fKD.sort = False
        fKD.rearrange = True
        create_tree(0)

    def _build_sort_array(self):
        """
        When running on a unigrid simulation, the kD tree isn't necessary.
        But we need to ensure that the points are sorted in the usual manner
        allowing values to be found via array indices.
        """
        mylog.info("Unigrid: finding cell centers.")
        xp = self.ds["x"]
        yp = self.ds["y"]
        zp = self.ds["z"]
        self.sizes = [np.unique(xp).size, np.unique(yp).size, np.unique(zp).size]
        self.sort = np.lexsort([zp, yp, xp])
        del xp, yp, zp
        self.ds.clear_data()

    def _build_fields_vals(self):
        """
        Builds an array to store the field values array.
        """
        self.fields_vals = np.empty((self.comm_size, len(self.fields)*2), \
            dtype='float64')
        # At the same time build a dict to label the columns.
        self.fields_columns = {}
        for i,field in enumerate(self.fields):
            self.fields_columns[field] = i

    def _build_points_array(self):
        """
        Initializes the array that contains the random points as all negatives
        to start with.
        """
        self.points = np.ones((self.comm_size, 6), dtype='float64') * -1.0

    def _setup_done_hooks_on_root(self):
        """
        Opens non-blocking receives on root pointing to all the other tasks
        """
        if self.mine != 0:
            return
        self.recv_done = {}
        for task in range(self.size):
            if task == self.mine: continue
            self.recv_done[task] = np.zeros(1, dtype='int64')
            self.done_hooks.append(self.comm.mpi_nonblocking_recv(self.recv_done[task], \
                task, tag=15))

    def _send_done_to_root(self):
        """
        Tell the root process that I'm done.
        """
        # If I've already done this, don't do it again.
        if self.sent_done: return
        if self.mine !=0:
            # I send when I *think* things should finish.
            self.send_done = np.ones(1, dtype='int64') * \
                (self.size / self.vol_ratio -1) + self.comm_cycle_count
            self.done_hooks.append(self.comm.mpi_nonblocking_send(self.send_done, \
                    0, tag=15))
        else:
            # As root, I need to mark myself!
            self.recv_done[0] = np.ones(1, dtype='int64') * \
                (self.size / self.vol_ratio -1) + self.comm_cycle_count
        self.sent_done = True

    def _should_cycle(self):
        """
        Determine if I should continue cycling the communication.
        """
        if self.mine == 0:
            # If other tasks aren't finished, this will return False.
            status = self.comm.mpi_Request_Testall(self.done_hooks)
            # Convolve this with with root's status.
            status = status * (self.generated_points == self.total_values)
            if status == 1:
                # If they are all finished, meaning Testall returns True,
                # and root has made its points, we find
                # the biggest value in self.recv_done and stop there.
                status = max(self.recv_done.values())
        else:
            status = 0
        # Broadcast the status from root - we stop only if root thinks we should
        # stop.
        status = self.comm.mpi_bcast(status)
        if status == 0: return True
        if self.comm_cycle_count < status:
            return True
        # If we've come this far, we're done.
        return False

    def _setup_recv_arrays(self):
        """
        Creates the recv buffers and calls a non-blocking MPI receive pointing
        to the left-hand neighbor.
        """
        self.recv_points = np.ones((self.comm_size, 6), dtype='float64') * -1.
        self.recv_fields_vals = np.zeros((self.comm_size, len(self.fields)*2), \
            dtype='float64')
        self.recv_gen_array = np.zeros(self.size, dtype='int64')
        self.recv_hooks.append(self.comm.mpi_nonblocking_recv(self.recv_points, \
            (self.mine-1)%self.size, tag=10))
        self.recv_hooks.append(self.comm.mpi_nonblocking_recv(self.recv_fields_vals, \
            (self.mine-1)%self.size, tag=20))
        self.recv_hooks.append(self.comm.mpi_nonblocking_recv(self.recv_gen_array, \
            (self.mine-1)%self.size, tag=40))

    def _send_arrays(self):
        """
        Send the data arrays to the right-hand neighbor.
        """
        self.send_hooks.append(self.comm.mpi_nonblocking_send(self.points,\
            (self.mine+1)%self.size, tag=10))
        self.send_hooks.append(self.comm.mpi_nonblocking_send(self.fields_vals,\
            (self.mine+1)%self.size, tag=20))
        self.send_hooks.append(self.comm.mpi_nonblocking_send(self.gen_array, \
            (self.mine+1)%self.size, tag=40))

    def _allsum_bin_hits(self):
        """
        Add up the hits to all the bins globally for all functions.
        """
        for fset in self._fsets:
            fset.too_low = self.comm.mpi_allreduce(fset.too_low, op='sum')
            fset.too_high = self.comm.mpi_allreduce(fset.too_high, op='sum')
            fset.binned = {}
            if self.mine == 0:
                mylog.info("Function %s had values out of range for these fields:" % \
                    fset.function.__name__)
                for i,field in enumerate(fset.out_labels):
                    mylog.info("Field %s had %d values too high and %d too low that were not binned." % \
                    (field, fset.too_high[i], fset.too_low[i]))
            for length in self.lengths:
                fset.length_bin_hits[length] = \
                    self.comm.mpi_allreduce(fset.length_bin_hits[length], op='sum')
                # Find out how many were successfully binned.
                fset.binned[length] = fset.length_bin_hits[length].sum()
                # Normalize the counts.
                fset.length_bin_hits[length] = \
                    fset.length_bin_hits[length].astype('float64') / \
                    fset.binned[length]
                # Return it to its original shape.
                fset.length_bin_hits[length] = \
                    fset.length_bin_hits[length].reshape(fset.bin_number)

    def _pick_random_points(self, length, size):
        """
        Picks out size random pairs separated by length *length*.
        """
        # First make random points inside this subvolume.
        r1 = np.empty((size,3), dtype='float64')
        for dim in range(3):
            r1[:,dim] = self.mt.uniform(low=self.ds.left_edge[dim],
                high=self.ds.right_edge[dim], size=size)
        # Next we find the second point, determined by a random
        # theta, phi angle. See Eqns. 1 & 2 from
        # http://mathworld.wolfram.com/SpherePointPicking.html,
        # but phi and theta are switched to the Physics convention.
        if self.constant_phi is None:
            phi = self.mt.uniform(low=0, high=2.*math.pi, size=size)
        else: phi = self.constant_phi * np.ones(size, dtype='float64')
        if self.constant_theta is None:
            v = self.mt.uniform(low=0., high=1, size=size)
            theta = np.arccos(2 * v - 1)
        else: theta = self.constant_theta * np.ones(size, dtype='float64')
        r2 = np.empty((size,3), dtype='float64')
        r2[:,0] = r1[:,0] + length * np.cos(phi) * np.sin(theta)
        r2[:,1] = r1[:,1] + length * np.sin(phi) * np.sin(theta)
        r2[:,2] = r1[:,2] + length * np.cos(theta)
        # Reflect so it's inside the (full) volume.
        r2 %= self.period
        return (r1, r2)

    def _find_nearest_cell(self, points):
        """
        Finds the closest grid cell for each point in a vectorized manner.
        """
        if self.nlevels == 0:
            pos = (points - self.ds.left_edge) / self.width
            n = (self.sizes[2] * pos[:,2]).astype('int32')
            n += self.sizes[2] * (self.sizes[1] * pos[:,1]).astype('int32')
            n += self.sizes[2] * self.sizes[1] * (self.sizes[0] * pos[:,0]).astype('int32')
        else:
            # Normalize the points to a 1-period for use only within the kdtree.
            points[:, 0] = points[:, 0] / self.period[0]
            points[:, 1] = points[:, 1] / self.period[1]
            points[:, 2] = points[:, 2] / self.period[2]
            fKD.qv_many = points.T
            fKD.nn_tags = np.asfortranarray(np.empty((1, points.shape[0]), dtype='int64'))
            fKD.find_many_nn_nearest_neighbors()
            # The -1 is for fortran counting.
            n = fKD.nn_tags[0,:] - 1
        return n

    def _get_fields_vals(self, points):
        """
        Given points, return the values for the fields we need for those
        points.
        """
        # First find the grid data index field.
        indices = self._find_nearest_cell(points)
        results = np.empty((len(indices), len(self.fields)), dtype='float64')
        # Put the field values into the columns of results.
        for field in self.fields:
            col = self.fields_columns[field]
            results[:, col] = self.stored_fields[field][indices]
        return results

    def _eval_points(self, length):
        # We need to loop over the points array at least once. Further
        # iterations only happen if we have added new points to the array,
        # but not as many as we want to, so we need to check again to see if
        # we can put more points into the buffer.
        added_points = True
        while added_points:
            # If we need to, add more points to the points array.
            if self.generated_points < self.total_values:
                # Look for 'empty' slots to put in new pairs.
                select = (self.points[:,0] < 0)
                ssum = select.sum()
                # We'll generate only as many points as we need to/can.
                size = min(ssum, self.total_values - self.generated_points)
                (new_r1,new_r2) = self._pick_random_points(length, size)
                self.generated_points += size
                # If size != select.sum(), we need to pad the end of new_r1/r2
                # which is what is effectively happening below.
                newpoints = np.ones((ssum, 6), dtype='float64') * -1.
                newpoints[:size,:3] = new_r1
                newpoints[:size,3:] = new_r2
                # Now we insert them into self.points.
                self.points[select] = newpoints
            else:
                added_points = False

            # If we have an empty buffer here, we can skip everything below.
            if (self.points < 0).all():
                added_points = False # Not strictly required, but clearer.
                break

            # Now we have a points array that is either full of unevaluated points,
            # or I don't need to make any new points and I'm just processing the
            # array. Start by finding the indices of the points I own.
            self.points.shape = (self.comm_size*2, 3) # Doesn't make a copy - fast!
            select = np.bitwise_or((self.points < self.ds.left_edge).any(axis=1),
                (self.points >= self.ds.right_edge).any(axis=1))
            select = np.invert(select)
            mypoints = self.points[select]
            if mypoints.size > 0:
                # Get the fields values.
                results = self._get_fields_vals(mypoints)
                # Put this into self.fields_vals.
                self.fields_vals.shape = (self.comm_size*2, len(self.fields))
                self.fields_vals[select] = results

            # Put our arrays back into their original shapes cheaply!
            if mypoints.size > 0:
                self.fields_vals.shape = (self.comm_size, len(self.fields)*2)
            self.points.shape = (self.comm_size, 6)

            # To run the functions, what is key is that the
            # second point in the pair is ours.
            second_points = self.points[:,3:]
            select = np.bitwise_or((second_points < self.ds.left_edge).any(axis=1),
                (second_points >= self.ds.right_edge).any(axis=1))
            select = np.invert(select)
            if select.any():
                points_to_eval = self.points[select]
                fields_to_eval = self.fields_vals[select]

                # Find the normal vector between our points.
                vec = np.abs(points_to_eval[:,:3] - points_to_eval[:,3:])
                norm = np.sqrt(np.sum(np.multiply(vec,vec), axis=1))
                # I wish there was a better way to do this, but I can't find it.
                for i, n in enumerate(norm):
                    vec[i] = np.divide(vec[i], n)

                # Now evaluate the functions.
                for fcn_set in self._fsets:
                    fcn_results = fcn_set._eval_st_fcn(fields_to_eval,points_to_eval,
                        vec)
                    fcn_set._bin_results(length, fcn_results)

                # Now clear the buffers at the processed points.
                self.points[select] = np.array([-1.]*6, dtype='float64')

            else:
                # We didn't clear any points, so we should move on with our
                # lives and pass this buffer along.
                added_points = False

    @parallel_blocking_call
    def write_out_means(self, fn = "%s.txt"):
        r"""Writes out the weighted-average value for each function for
        each dimension for each ruler length to a text file. The data is written
        to files of the name 'function_name.txt' in the current working
        directory.

        Examples
        --------
        >>> tpf.write_out_means()
        """
        for fset in self._fsets:
            fp = self.comm.write_on_root(fn % fset.function.__name__)
            fset._avg_bin_hits()
            line = "# length".ljust(sep)
            line += "count".ljust(sep)
            for dim in fset.dims:
                line += ("%s" % fset.out_labels[dim]).ljust(sep)
            fp.write(line + "\n")
            for length in self.lengths:
                line = ("%1.5e" % length).ljust(sep)
                line += ("%d" % fset.binned[length]).ljust(sep)
                for dim in fset.dims:
                    if fset.sqrt[dim]:
                        line += ("%1.5e" % \
                            math.sqrt(fset.length_avgs[length][dim])).ljust(sep)
                    else:
                        line += ("%1.5e" % \
                            fset.length_avgs[length][dim]).ljust(sep)
                line += "\n"
                fp.write(line)
            fp.close()

    @parallel_root_only
    def write_out_arrays(self, fn = "%s.h5"):
        r"""Writes out the raw probability bins and the bin edges to an HDF5 file
        for each of the functions. The files are named
        'function_name.txt' and saved in the current working directory.

        Examples
        --------
        >>> tpf.write_out_arrays()
        """
        if self.mine == 0:
            for fset in self._fsets:
                f = h5py.File(fn % fset.function.__name__, "w")
                bin_names = []
                prob_names = []
                bin_counts = []
                for dim in fset.dims:
                    f.create_dataset("/bin_edges_%02d_%s" % \
                        (dim, fset.out_labels[dim]), \
                        data=fset.bin_edges[dim])
                    bin_names.append("/bin_edges_%02d_%s" % \
                    (dim, fset.out_labels[dim]))
                for i,length in enumerate(self.lengths):
                    f.create_dataset("/prob_bins_%05d" % i, \
                        data=fset.length_bin_hits[length])
                    prob_names.append("/prob_bins_%05d" % i)
                    bin_counts.append([fset.too_low.sum(), fset.binned[length],
                        fset.too_high.sum()])
                f.create_dataset("/bin_edges_names", data=bin_names)
                #f.create_dataset("/prob_names", data=prob_names)
                f.create_dataset("/lengths", data=self.lengths)
                f.create_dataset("/counts", data=bin_counts)
                f.close()

    @parallel_root_only
    def write_out_correlation(self):
        r"""A special output function for doing two point correlation functions.
        Outputs the correlation function xi(r) in a text file
        'function_name_corr.txt' in the current working directory.

        Examples
        --------
        >>> tpf.write_out_correlation()
        """
        for fset in self._fsets:
            # Only operate on correlation functions.
            if fset.corr_norm is None: continue
            fp = self.comm.write_on_root("%s_correlation.txt" % fset.function.__name__)
            line = "# length".ljust(sep)
            line += "\\xi".ljust(sep)
            fp.write(line + "\n")
            xi = fset._corr_sum_norm()
            for length in self.lengths:
                line = ("%1.5e" % length).ljust(sep)
                line += ("%1.5e" % xi[length]).ljust(sep)
                fp.write(line + "\n")
            fp.close()

class FcnSet(TwoPointFunctions):
    def __init__(self,tpf, function, min_edge, out_labels, sqrt, corr_norm):
        self.tpf = tpf # The overarching TPF class
        self.function = function # Function to eval between the two points.
        self.min_edge = min_edge # The length of the minimum edge of the box.
        self.out_labels = out_labels # For output.
        self.sqrt = sqrt # which columns to sqrt on output.
        self.corr_norm = corr_norm # A number used to normalize a correlation function.
        # These below are used to track how many times the function returns
        # unbinned results.
        self.too_low = np.zeros(len(self.out_labels), dtype='int32')
        self.too_high = np.zeros(len(self.out_labels), dtype='int32')

    def set_pdf_params(self, bin_type="lin", bin_number=1000, bin_range=None):
        r"""Set the parameters used to build the Probability Distribution Function
        for each ruler length for this function. The values output by the
        function are slotted into the bins described here.

        Parameters
        ----------
        bin_type : String
            Controls the edges of the bins spaced evenly in
            logarithmic or linear space, set by "log" or "lin", respectively.
            A single string, or list of strings for N-dim binning.
            Default = "lin".
        bin_number : Integer
            Sets how many bins to create, evenly spaced by the above
            parameter. A single integer, or a list of integers for N-dim
            binning. Default = 1000.
        bin_range : Float
            A pair of values giving the range for the bins.
            A pair of floats (a list), or a list of pairs for N-dim binning.
            Default = None.

        Examples
        --------
        >>> f1.set_pdf_params(bin_type='log', bin_range=[5e4, 5.5e13],
        ... bin_number=1000)
        """
        # This should be called after setSearchParams.
        if not hasattr(self.tpf, "lengths"):
            mylog.error("Please call setSearchParams() before calling setPDFParams().")
            return None
        # Make sure they're either all lists or only one is.
        input = [bin_type, bin_number, bin_range]
        lists = 0
        for thing in input:
            if type(thing) == list:
                lists += 1
        if lists > 1 and lists < 3:
            mylog.error("Either all the inputs need to be lists, or only one.")
            return None
        # Make sure they're all the same length if they're lists.
        if lists == 3:
            first_len = 0
            for thing in input:
                if first_len == 0:
                    first_len = len(thing)
                    if first_len == 0:
                        mylog.error("Input cannot be an empty list.")
                        return None
                    continue
                if first_len != len(thing):
                    mylog.error("All the inputs need to have the same length.")
                    return None
        # If they are not all lists, put the input into lists for convenience.
        if lists == 1:
            bin_type, bin_number = [bin_type], [bin_number]
            bin_range = [bin_range]
        self.bin_type = bin_type
        self.bin_number = np.array(bin_number) - 1
        self.dims = range(len(bin_type))
        # Create the dict that stores the arrays to store the bin hits, and
        # the arrays themselves.
        self.length_bin_hits = {}
        for length in self.tpf.lengths:
            # It's easier to index flattened, but will be unflattened later.
            self.length_bin_hits[length] = np.zeros(self.bin_number,
                dtype='int64').flatten()
        # Create the bin edges for each dimension.
        # self.bins is indexed by dimension
        self.bin_edges = {}
        for dim in self.dims:
            # Error check.
            if len(bin_range[dim]) != 2:
                raise ValueError("bin_range must have two values.")
            if bin_range[dim][1] <= bin_range[dim][0]:
                raise ValueError("bin_range[1] must be larger than bin_range[0]")
            # Make the edges for this dimension.
            if bin_type[dim] == "lin":
                self.bin_edges[dim] = np.linspace(bin_range[dim][0], bin_range[dim][1],
                    bin_number[dim])
            elif bin_type[dim] == "log":
                self.bin_edges[dim] = np.logspace(math.log10(bin_range[dim][0]),
                    math.log10(bin_range[dim][1]), bin_number[dim])
            else:
                raise SyntaxError("bin_edges is either \"lin\" or \"log\".")

    def _eval_st_fcn(self, results, points, vec):
        """
        Return the value of the function using the provided results.
        """
        return self.function(results[:,:len(self.tpf.fields)],
            results[:,len(self.tpf.fields):], points[:,:3], points[:,3:], vec)
        """
        NOTE - A function looks like:
        def stuff(a,b,r1,r2, vec):
            return [(a[0] - b[0])/(a[1] + b[1])]
        where a and b refer to different points in space and the indices
        are for the different fields, which are given when the function is
        added. The results need to be a list or array even if it's only one
        item.
        """

    def _bin_results(self, length, results):
        """
        Add hits to the bins corresponding to these results. length_hit_bins
        is flattened, so we need to figure out the offset for this hit by
        factoring the sizes of the other dimensions.
        """
        hit_bin = np.zeros(results.shape[0], dtype='int64')
        multi = 1
        good = np.ones(results.shape[0], dtype='bool')
        for dim in range(len(self.out_labels)):
            for d1 in range(dim):
                multi *= self.bin_edges[d1].size
            if dim == 0 and len(self.out_labels)==1:
                try:
                    digi = np.digitize(results, self.bin_edges[dim])
                except ValueError:
                    # The user probably did something like
                    # return a * b rather than
                    # return a[0] * b[0], which will only happen
                    # for single field functions.
                    digi = np.digitize(results[0], self.bin_edges[dim])
            else:
                digi = np.digitize(results[:,dim], self.bin_edges[dim])
            too_low = (digi == 0)
            too_high = (digi == self.bin_edges[dim].size)
            self.too_low[dim] += (too_low).sum()
            self.too_high[dim] += (too_high).sum()
            newgood = np.bitwise_and(np.invert(too_low), np.invert(too_high))
            good = np.bitwise_and(good, newgood)
            hit_bin += np.multiply((digi - 1), multi)
        digi_bins = np.arange(self.length_bin_hits[length].size+1)
        hist, digi_bins = np.histogram(hit_bin[good], digi_bins)
        self.length_bin_hits[length] += hist

    def _dim_sum(self, a, dim):
        """
        Given a multidimensional array a, this finds the sum over all the
        elements leaving the dimension dim untouched.
        """
        dims = np.arange(len(a.shape))
        dims = np.flipud(dims)
        gt_dims = dims[dims > dim]
        lt_dims = dims[dims < dim]
        iter_dims = np.concatenate((gt_dims, lt_dims))
        for this_dim in iter_dims:
            a = a.sum(axis=this_dim)
        return a

    def _avg_bin_hits(self):
        """
        For each dimension and length of bin_hits return the weighted average.
        """
        self.length_avgs = defaultdict(dict)
        for length in self.tpf.lengths:
            for dim in self.dims:
                self.length_avgs[length][dim] = \
                    (self._dim_sum(self.length_bin_hits[length], dim) * \
                    ((self.bin_edges[dim][:-1] + self.bin_edges[dim][1:]) / 2.)).sum()

    def _corr_sum_norm(self):
        """
        Return the correlations xi for this function. We are tacitly assuming
        that all correlation functions are one dimensional.
        """
        xi = {}
        for length in self.tpf.lengths:
            xi[length] = -1 + np.sum(self.length_bin_hits[length] * \
                self.bin_edges[0][:-1]) / self.corr_norm
        return xi

ext: py

reason: import numpy


text:
"""
Decodes events like sound stimulation or luminance/weight conditions.
@author: bartulem
"""

import os
import gc
import time
import numpy as np
from numba import njit
from tqdm import tqdm
from scipy.stats import zscore
import neural_activity
import select_clusters
import sessions2load

@njit(parallel=False)
def correlate_quickly(big_x, big_x_mean, big_y, big_y_mean):
    big_x_shape_0 = big_x.shape[0]
    big_y_shape_0 = big_y.shape[0]
    big_y_shape_1 = big_y.shape[1]
    big_x_mean = np.reshape(big_x_mean, (big_x_shape_0, 1))
    big_y_mean = np.reshape(big_y_mean, (big_y_shape_0, 1))
    x_ds = big_x - big_x_mean
    y_ds = big_y - big_y_mean
    r_num = (x_ds.reshape((big_x_shape_0, 1, big_y_shape_1)) * y_ds).sum(axis=2)
    r_den = np.sqrt(((x_ds ** 2).sum(axis=1).reshape((big_x_shape_0, 1))) * ((y_ds ** 2).sum(axis=1).reshape((1, big_y_shape_0))))
    r = r_num / r_den
    nan_positions = np.isnan(r)
    return r, nan_positions

def predict_events(pred_arr_len, fold_num, train_folds, test_folds,
                   activity_arr, event_arr, fe, three_sessions=False):
    pred_events = np.zeros(pred_arr_len).astype(np.float32)
    for fold_idx in range(fold_num):
        training_frames = train_folds[fold_idx]
        test_frames = test_folds[fold_idx]
        train_arr = activity_arr.take(indices=training_frames, axis=0).astype(np.float32)
        test_arr = activity_arr.take(indices=test_frames, axis=0).astype(np.float32)
        corr_arr, nan_pos = correlate_quickly(big_x=train_arr,
                                              big_x_mean=train_arr.mean(axis=1),
                                              big_y=test_arr,
                                              big_y_mean=test_arr.mean(axis=1))
        # exchange nans for lowest value
        corr_arr[nan_pos] = -1

        # find the best congruent train frame
        # since argmax would always take the first frame if multiple frames had the same max value,
        # it's necessary to make a random selection from all the frames that have the max value
        max_corr_each_frame = np.nanmax(corr_arr, axis=0).astype(np.float32)
        max_corr_train_frames_raw = np.zeros(max_corr_each_frame.shape[0]).astype(np.int32)
        for test_fr in range(max_corr_each_frame.shape[0]):
            max_corr_train_frames_raw[test_fr] = np.random.choice(np.where(corr_arr[:, test_fr] == max_corr_each_frame[test_fr])[0]).astype(np.int32)

        actual_train_frames = training_frames.take(max_corr_train_frames_raw)

        # get event values
        if not three_sessions:
            pred_events[fe[fold_idx]:fe[fold_idx + 1]] = event_arr.take(actual_train_frames)
        else:
            pred_events = event_arr.take(actual_train_frames)

    return pred_events

def choose_012_clusters(the_input_012, cl_gr_dir, sp_prof_csv, cl_areas, cl_type, dec_type, desired_profiles):
    chosen_clusters = []
    extra_chosen_clusters = {0: [], 1: []}
    cluster_dict = {0: [], 1: [], 2: []}
    for session_idx, one_session in enumerate(the_input_012):
        cluster_dict[session_idx] = select_clusters.ClusterFinder(session=one_session,
                                                                  cluster_groups_dir=cl_gr_dir,
                                                                  sp_profiles_csv=sp_prof_csv).get_desired_clusters(filter_by_area=cl_areas,
                                                                                filter_by_cluster_type=cl_type,
                                                                                filter_by_spiking_profile=desired_profiles)

    # find clusters present in all 3 sessions
    for one_cl in list(set(cluster_dict[0]).intersection(cluster_dict[1], cluster_dict[2])):
        chosen_clusters.append(one_cl)

    if dec_type == 'luminance':
        # find cells present only in the dark session
        if 'V' in cl_areas:
            for d_cl in cluster_dict[1]:
                if d_cl not in cluster_dict[0] and d_cl not in cluster_dict[2]:
                    extra_chosen_clusters[1].append(d_cl)

        # find clusters present in two light session but not in the dark
        if 'V' in cl_areas:
            for l_cl in cluster_dict[0]:
                if l_cl in cluster_dict[2] and l_cl not in cluster_dict[1]:
                    extra_chosen_clusters[0].append(l_cl)

    # all chosen clusters
    all_clusters = chosen_clusters + extra_chosen_clusters[0] + extra_chosen_clusters[1]

    return all_clusters, chosen_clusters, extra_chosen_clusters, cluster_dict

class Decoder:

    def __init__(self, input_file='', save_results_dir='', input_012=None,
                 cluster_groups_dir='/.../cluster_groups_info',
                 sp_profiles_csv='/.../spiking_profiles.csv',
                 number_of_decoding_per_run=10, decoding_cell_number_array=np.array([5, 10, 20, 50, 100]),
                 fold_n=3, shuffle_num=1000,
                 to_smooth=False, smooth_sd=1, smooth_axis=0, condense=True,
                 cluster_areas=None, cluster_type=True, animal_names=None,
                 profiles_to_include=True, to_zscore=False):
        if cluster_areas is None:
            cluster_areas = ['A']
        if input_012 is None:
            input_012 = ['', '', '']
        if animal_names is None:
            animal_names = ['kavorka', 'frank', 'johnjohn', 'roy', 'jacopo', 'bruno', 'crazyjoe']
        self.input_file = input_file
        self.save_results_dir = save_results_dir
        self.input_012 = input_012
        self.cluster_groups_dir = cluster_groups_dir
        self.sp_profiles_csv = sp_profiles_csv
        self.number_of_decoding_per_run = number_of_decoding_per_run
        self.decoding_cell_number_array = decoding_cell_number_array
        self.fold_n = fold_n
        self.shuffle_num = shuffle_num
        self.to_smooth = to_smooth
        self.smooth_sd = smooth_sd
        self.smooth_axis = smooth_axis
        self.condense = condense
        self.cluster_areas = cluster_areas
        self.cluster_type = cluster_type
        self.animal_names = animal_names
        self.profiles_to_include = profiles_to_include
        self.to_zscore = to_zscore

    def decode_sound_stim(self, **kwargs):
        """
        Description
        ----------
        This method uses a simple nearest neighbor decoder to predict sound stimulation
        events. More specifically, we bin the spike train of every cluster in 100 ms bins
        and smooth it with a 3 bin Gaussian kernel. Since our three animals have a varying
        number of auditory single units (K=288, JJ=132, F=198), for each one we choose a
        distinct combination of units (either 5, 10, 20, 50 or 100) to decode the presence
        of the sound stimulus in each run (10 runs in total). In each run, we divided the data
        in three folds where 1/3 of the data was the test set and 2/3 were the training set.
        For each test set population vector we computed Pearson correlations to every population
        vector in the training set. We obtained a predicted sound stimulus value for each test
        frame by assigning it the sound stimulus value of the most correlated training set population
        vector. Decoding accuracy was defined as the proportion of correctly matched stimulus states
        across the entire recording session. We also shuffled the spikes of these units 1000 times
        in the first run to obtain the null-distribution of decoded accuracy.
        ----------

        Parameters
        ----------
        cluster_groups_dir (str)
            The directory with the cluster_groups .json file; defaults to '/.../cluster_groups_info'.
        sp_profiles_csv (str)
            The sp_profiles .csv file; defaults to '/.../spiking_profiles.csv'.
        cluster_areas (list)
            Cluster area(s) of choice; defaults to ['A'].
        cluster_type (str)
            Cluster type of choice; defaults to True.
        number_of_decoding_per_run (int)
            The number of times to decode event per run; defaults to 10.
        decoding_cell_number_array (np.ndarray):
            An array of numbers of cells to decode with; defaults to np.array([5, 10, 20, 50, 100]).
        shuffle_num (int)
            Number of shuffles; defaults to 1000.
        to_smooth (bool)
            Smooth spike trains; defaults to False.
        smooth_sd (int)
            The SD of the smoothing window; defaults to 1 (bin).
        smooth_axis (int)
            The smoothing axis in a 2D array; defaults to 0.
        fold_n (int)
            The number of folds for decoding; defaults to 3.
        condense (bool)
            Yey or ney on the spike array condensing; defaults to True.
        ----------

        Returns
        ----------
        sound_decoding_accuracy (.npy file)
            The decoding accuracy data for a given input file.
        sound_shuffled_decoding_accuracy (.npy file)
            The shuffled data decoding accuracy for a given input file.
        ----------
        """

        # keep time
        start_time = time.time()

        # get animal name
        animal_name = [name for name in self.animal_names if name in self.input_file][0]

        # choose clusters you'd like to decode with
        chosen_clusters = select_clusters.ClusterFinder(session=self.input_file,
                                                        cluster_groups_dir=self.cluster_groups_dir,
                                                        sp_profiles_csv=self.sp_profiles_csv).get_desired_clusters(filter_by_area=self.cluster_areas,
                                                                                filter_by_cluster_type=self.cluster_type)
        # get total frame count and sound array
        file_name, extracted_frame_info = sessions2load.Session(session=self.input_file).data_loader(extract_variables=['total_frame_num', 'imu_sound'])

        # get activity dictionary
        file_id, activity_dictionary, purged_spikes_dict = neural_activity.Spikes(input_file=self.input_file).convert_activity_to_frames_with_shuffles(get_clusters=chosen_clusters,
                                                                                to_shuffle=True,
                                                                                condense_arr=self.condense)
        if self.condense:
            sound_array = neural_activity.condense_frame_arrays(frame_array=extracted_frame_info['imu_sound'], arr_type=False)
            total_frame_num = sound_array.shape[0]
        else:
            sound_array = extracted_frame_info['imu_sound']
            total_frame_num = extracted_frame_info['total_frame_num']

        # get fold edges
        fold_edges = np.floor(np.linspace(0, total_frame_num, self.fold_n + 1)).astype(np.int32)

        # get train / test indices for each fold
        train_indices_for_folds = []
        test_indices_for_folds = []
        for fold in range(self.fold_n):
            all_frames = np.arange(0, total_frame_num)
            one_fold_arr = np.arange(fold_edges[fold], fold_edges[fold + 1])
            test_indices_for_folds.append(one_fold_arr.astype(np.int32))
            train_indices_for_folds.append(np.setdiff1d(all_frames, one_fold_arr).astype(np.int32))

        # conduct decoding
        decoding_accuracy = np.zeros((self.decoding_cell_number_array.shape[0], self.number_of_decoding_per_run))
        shuffled_decoding_accuracy = np.zeros((self.decoding_cell_number_array.shape[0], self.shuffle_num))
        for decode_num in range(self.number_of_decoding_per_run):
            for ca_idx, cell_amount in enumerate(self.decoding_cell_number_array):
                clusters_array = np.zeros((total_frame_num, cell_amount)).astype(np.float32)

                if decode_num == 0:
                    shuffled_clusters_array = np.zeros((self.shuffle_num, total_frame_num, cell_amount))

                # shuffle cluster names
                np.random.shuffle(chosen_clusters)

                # select clusters
                selected_clusters = chosen_clusters[:cell_amount]

                # get all cell / shuffled data in their respective arrays
                for sc_idx, selected_cluster in enumerate(selected_clusters):
                    clusters_array[:, sc_idx] = activity_dictionary[selected_cluster]['activity'].todense().astype(np.float32).copy()
                    if decode_num == 0:
                        for shuffle_idx in range(self.shuffle_num):
                            shuffled_clusters_array[shuffle_idx, :, sc_idx] = activity_dictionary[selected_cluster]['shuffled'][shuffle_idx].todense().astype(np.float32).copy()

                # smooth spike trains if desired
                if self.to_smooth:
                    clusters_array = neural_activity.gaussian_smoothing(array=clusters_array, sigma=self.smooth_sd, axis=self.smooth_axis).astype(np.float32)
                    if decode_num == 0:
                        shuffled_clusters_array = neural_activity.gaussian_smoothing(array=shuffled_clusters_array, sigma=self.smooth_sd, axis=self.smooth_axis + 1).astype(np.float32)

                # go through folds and predict sound
                predicted_sound_events = predict_events(pred_arr_len=total_frame_num, fold_num=self.fold_n, train_folds=train_indices_for_folds,
                                                        test_folds=test_indices_for_folds, activity_arr=clusters_array, event_arr=sound_array,
                                                        fe=fold_edges)
                if decode_num == 0:
                    shuffle_predicted_sound_events = np.zeros((total_frame_num, self.shuffle_num))
                    for sh in range(self.shuffle_num):
                        shuffle_predicted_sound_events[:, sh] = predict_events(pred_arr_len=total_frame_num, fold_num=self.fold_n, train_folds=train_indices_for_folds,
                                                                               test_folds=test_indices_for_folds, activity_arr=shuffled_clusters_array[sh], event_arr=sound_array,
                                                                               fe=fold_edges)

                # calculate accuracy and fill in the array
                decoding_accuracy[ca_idx, decode_num] = ((predicted_sound_events - sound_array) == 0).sum() / predicted_sound_events.shape[0]
                if decode_num == 0:
                    for sh_idx in range(self.shuffle_num):
                        shuffled_decoding_accuracy[ca_idx, sh_idx] = ((shuffle_predicted_sound_events[:, sh_idx] - sound_array) == 0).sum() \
                                                                     / shuffle_predicted_sound_events.shape[0]
                # free memory
                gc.collect()

        # save results as .npy files
        np.save(f'{self.save_results_dir}{os.sep}{animal_name}_sound_decoding_accuracy_{self.cluster_areas[0]}_clusters', decoding_accuracy)
        np.save(f'{self.save_results_dir}{os.sep}{animal_name}_sound_shuffled_decoding_accuracy_{self.cluster_areas[0]}_clusters', shuffled_decoding_accuracy)

        print("Sound decoding complete! It took {:.2f} hours.".format((time.time() - start_time) / 3600))

    def decode_session_type(self, **kwargs):
        """
        Description
        ----------
        This method uses a SVM/nearest neighbor decoder to predict the luminance / weight condition
        of individual timepoints in a separate recording session (i.e. the second light session).
        We take the last quarter of recording of condition 1 (e.g. light), the first quarter of
        recording of condition 2 (e.g. dark) and the second half of a separate recording of condition
        1 (e.g. second light session). The spike trains of every cluster common to all three sessions
        are binned in 100 ms bins (and smoothed with a 3 bin Gaussian). Since our three animals have
        a varying number of visual and  A1 single units, we choose a distinct combination of units
        (either 5, 10, 20, 35 or 50) to decode the condition in each run (100 runs in total).
        For each test set population vector (e.g. second half of the second light session), we computed
        Pearson correlations to every population vector in the training set (light 1 + dark). We obtained
        a predicted condition status for each test frame by assigning it the condition status of the most
        correlated training set population vector. Decoding accuracy was defined as the proportion of
        correctly matched stimulus states across the test session. Since the luminance condition didn't
        change within a session, shuffling spike trains would not make sense because even time shifted
        activity, if it's overall lower/higher relative to the other session, would still enable accurate
        decoding. Instead, we randomly permuted the joint unit activity (quarter light / quarter dark) at each
        training time point a 1000 times in the first run to obtain the null-distribution of decoded accuracy.

        NB: To run this we used the following parameters:
            cluster_areas = ['V'] or ['A1']
            cluster_type = 'good'
            to_smooth = True
            smooth_sd = 1
            condense = True
            decoding_cell_number_array = np.array([5, 10, 20, 35, 50])
            number_of_decoding_per_run = 100
            pop_vector_decoding = True
        ----------

        Parameters
        ----------
        **kwargs (dictionary)
        condensed_bin_size (int)
            Condensed bin size; defaults to 100 (ms).
        decode_what (str)
            What are you decoding; defaults to 'luminance'.
        ----------

        Returns
        ----------
        decoding_accuracy (.npy file)
            The decoding accuracy data for a given input file pair.
        shuffled_decoding_accuracy (.npy file)
            The shuffled data decoding accuracy for a given input file pair.
        ----------
        """

        condensed_bin_size = kwargs['condensed_bin_size'] if 'condensed_bin_size' in kwargs.keys() and type(kwargs['condensed_bin_size']) == int else 100
        decode_what = kwargs['decode_what'] if 'decode_what' in kwargs.keys() and type(kwargs['decode_what']) == str else 'luminance'

        # keep time
        session_type_decoding_start_time = time.time()

        # get animal name
        animal_name = [name for name in self.animal_names if name in self.input_012[0]][0]

        # choose clusters for decoding
        all_clusters, chosen_clusters, extra_chosen_clusters, cluster_dict = choose_012_clusters(the_input_012=self.input_012,
                                                                                cl_gr_dir=self.cluster_groups_dir,
                                                                                sp_prof_csv=self.sp_profiles_csv,
                                                                                cl_areas=self.cluster_areas,
                                                                                cl_type=self.cluster_type,
                                                                                dec_type=decode_what,
                                                                                desired_profiles=self.profiles_to_include)

        # get total frame count in each session
        zero_ses_name, zero_extracted_frame_info = sessions2load.Session(session=self.input_012[0]).data_loader(extract_variables=['total_frame_num'])
        first_ses_name, first_extracted_frame_info = sessions2load.Session(session=self.input_012[1]).data_loader(extract_variables=['total_frame_num'])
        second_ses_name, second_extracted_frame_info = sessions2load.Session(session=self.input_012[2]).data_loader(extract_variables=['total_frame_num'])

        # get total frame number in future array
        if self.condense:
            total_frame_num = np.array([zero_extracted_frame_info['total_frame_num'],
                                        first_extracted_frame_info['total_frame_num'],
                                        second_extracted_frame_info['total_frame_num']]).min() // int(120. * (condensed_bin_size / 1e3))
        else:
            total_frame_num = np.array([zero_extracted_frame_info['total_frame_num'],
                                        first_extracted_frame_info['total_frame_num'],
                                        second_extracted_frame_info['total_frame_num']]).min()

        middle_change_point = total_frame_num // 2
        fold_edges = np.array([0, middle_change_point // 2, middle_change_point, total_frame_num])
        third_durations = np.diff(fold_edges)
        decoding_event_array = np.concatenate((np.ones(third_durations[0]), np.zeros(third_durations[1]), np.ones(third_durations[2]))).astype(np.float32)

        # get activity dictionary
        zero_first_second_activity = {0: {}, 1: {}, 2: {}}
        for file_idx, one_file in enumerate(self.input_012):
            if 'V' in self.cluster_areas and decode_what == 'luminance':
                if file_idx == 0 or file_idx == 2:
                    chosen_idx = 0
                else:
                    chosen_idx = 1
                file_id, activity_dictionary, purged_spikes_dict = neural_activity.Spikes(input_file=one_file).convert_activity_to_frames_with_shuffles(
                    get_clusters=chosen_clusters + extra_chosen_clusters[chosen_idx],
                    to_shuffle=False,
                    condense_arr=self.condense)
            else:
                file_id, activity_dictionary, purged_spikes_dict = neural_activity.Spikes(input_file=one_file).convert_activity_to_frames_with_shuffles(get_clusters=chosen_clusters,
                                                                                to_shuffle=False,
                                                                                condense_arr=self.condense)
            zero_first_second_activity[file_idx] = activity_dictionary

        # get train / test indices
        train_indices_for_folds = [np.arange(0, middle_change_point).astype(np.int32)]
        test_indices_for_folds = [np.arange(middle_change_point, total_frame_num).astype(np.int32)]

        # conduct decoding
        decoding_accuracy = np.zeros((self.decoding_cell_number_array.shape[0], self.number_of_decoding_per_run))
        shuffled_decoding_accuracy = np.zeros((self.decoding_cell_number_array.shape[0], self.shuffle_num))
        for decode_num in tqdm(range(self.number_of_decoding_per_run)):
            for ca_idx, cell_amount in enumerate(self.decoding_cell_number_array):
                clusters_array = np.zeros((total_frame_num, cell_amount)).astype(np.float32)

                if decode_num == 0:
                    shuffled_clusters_array = np.zeros((self.shuffle_num, total_frame_num, cell_amount)).astype(np.float32)

                # shuffle cluster names
                np.random.shuffle(all_clusters)

                # select clusters
                selected_clusters = all_clusters[:cell_amount]

                # get all cell / shuffled data in their respective arrays
                for sc_idx, selected_cluster in enumerate(selected_clusters):
                    for condition_type in range(len(zero_first_second_activity.keys())):
                        seq_len = third_durations[condition_type]
                        if selected_cluster in zero_first_second_activity[condition_type].keys():
                            if condition_type == 1:
                                temp_cl_arr = zero_first_second_activity[condition_type][selected_cluster]['activity'][:seq_len].todense().astype(np.float32).copy()
                            else:
                                temp_cl_arr = zero_first_second_activity[condition_type][selected_cluster]['activity'][-seq_len:].todense().astype(np.float32).copy()
                            if self.to_smooth:
                                temp_cl_arr = neural_activity.gaussian_smoothing(array=temp_cl_arr, sigma=self.smooth_sd, axis=self.smooth_axis).astype(np.float32)
                            if self.to_zscore:
                                temp_cl_arr = zscore(temp_cl_arr)
                            clusters_array[fold_edges[condition_type]: fold_edges[condition_type + 1], sc_idx] = temp_cl_arr

                # prepare array for shuffling
                if decode_num == 0:
                    copy_clu_train_arr = clusters_array[:middle_change_point, :].copy()
                    copy_clu_test_arr = clusters_array[middle_change_point:, :].copy()
                    for shuffle_idx in range(self.shuffle_num):
                        np.random.shuffle(copy_clu_train_arr)
                        shuffled_clusters_array[shuffle_idx, :, :] = np.concatenate((copy_clu_train_arr, copy_clu_test_arr))

                # go through folds and predict event
                predicted_condition_events = predict_events(pred_arr_len=total_frame_num - third_durations[2], fold_num=1, train_folds=train_indices_for_folds,
                                                            test_folds=test_indices_for_folds, activity_arr=clusters_array, event_arr=decoding_event_array.copy(),
                                                            fe=fold_edges, three_sessions=True)

                if decode_num == 0:
                    shuffle_predicted_condition_events = np.zeros((total_frame_num - middle_change_point, self.shuffle_num))
                    for sh in range(self.shuffle_num):
                        shuffle_predicted_condition_events[:, sh] = predict_events(pred_arr_len=total_frame_num - third_durations[2], fold_num=1, train_folds=train_indices_for_folds,
                                                                                test_folds=test_indices_for_folds, activity_arr=shuffled_clusters_array[sh],
                                                                                event_arr=decoding_event_array.copy(),
                                                                                fe=fold_edges, three_sessions=True)

                # calculate accuracy and fill in the array
                decoding_accuracy[ca_idx, decode_num] = ((predicted_condition_events - np.ones(total_frame_num - middle_change_point)) == 0).sum() / predicted_condition_events.shape[0]

                if decode_num == 0:
                    for sh_idx in range(self.shuffle_num):
                        shuffled_decoding_accuracy[ca_idx, sh_idx] = ((shuffle_predicted_condition_events[:, sh_idx] - np.ones(total_frame_num - middle_change_point)) == 0).sum() \
                                                                     / shuffle_predicted_condition_events[:, sh_idx].shape[0]

                # free memory
                gc.collect()

        # save results as .npy files
        np.save(f'{self.save_results_dir}{os.sep}{animal_name}_{decode_what}_decoding_accuracy_{self.cluster_areas[0]}_clusters', decoding_accuracy)
        np.save(f'{self.save_results_dir}{os.sep}{animal_name}_{decode_what}_shuffled_decoding_accuracy_{self.cluster_areas[0]}_clusters', shuffled_decoding_accuracy)

        print("Decoding complete! It took {:.2f} hours.".format((time.time() - session_type_decoding_start_time) / 3600))

max_stars_repo_name: bartulem/KISN-pancortical-kinematics

max_issues_repo_name: bartulem/KISN-pancortical-kinematics

max_forks_repo_name: bartulem/KISN-pancortical-kinematics

reason: import numpy,from scipy,from numba


text:
# -*- coding: utf-8 -*-
"""
registration.py contains a simple 2D image registration function,
which registers images by checking the individual horizontal and vertical shifts
"""
from __future__ import division, unicode_literals, print_function

from .signal import find_phase, vertical_mean, horizontal_mean

import numpy as np

def translation_2x1d(image_a=None, image_b=None, ffts_a=(), ffts_b=(), return_a=False, return_b=False):
    """

    :param image_a:
    :param image_b:
    :param ffts_a:
    :param ffts_b:
    :param return_a:
    :param return_b:
    :return:
    """
    if ffts_a != ():
        fft_av, fft_ah = ffts_a
        signal_av, signal_ah = None, None
    else:
        fft_av, fft_ah = None, None
        signal_av, signal_ah = vertical_mean(image_a), horizontal_mean(image_a)

    if ffts_b != ():
        fft_bv, fft_bh = ffts_b
        signal_bv, signal_bh = None, None
    else:
        fft_bv, fft_bh = None, None
        signal_bv, signal_bh = vertical_mean(image_b), horizontal_mean(image_b)

    v, fft_av, fft_bv = find_phase(signal_1=signal_av, signal_2=signal_bv,
                                   fft_1=fft_av, fft_2=fft_bv, return_1=True, return_2=True)

    h, fft_ah, fft_bh = find_phase(signal_1=signal_ah, signal_2=signal_bh,
                                   fft_1=fft_ah, fft_2=fft_bh, return_1=True, return_2=True)

    result = ([float(-v), float(-h)],)

    if return_a:
        result += ((fft_av, fft_ah),)

    if return_b:
        result += ((fft_bv, fft_bh),)

    return result

def shift_image(image, shift, background='input'):
    """

    :param image:
    :param shift:
    :param background:
    :return: :raise ValueError:
    """

    vertical, horizontal = shift
    vertical, horizontal = round(vertical), round(horizontal)
    height, width = image.shape

    if vertical < 0:
        source_vertical_lower = -vertical
        source_vertical_upper = height
        destination_vertical_lower = 0
        destination_vertical_upper = vertical
    else:
        source_vertical_lower = 0
        source_vertical_upper = height - vertical
        destination_vertical_lower = vertical
        destination_vertical_upper = height

    if horizontal < 0:
        source_horizontal_lower = -horizontal
        source_horizontal_upper = width
        destination_horizontal_lower = 0
        destination_horizontal_upper = horizontal
    else:
        source_horizontal_lower = 0
        source_horizontal_upper = width - horizontal
        destination_horizontal_lower = horizontal
        destination_horizontal_upper = width

    if background == 'input':
        new_image = image.copy()
    elif background == 'blank':
        new_image = np.zeros_like(image)
    else:
        raise ValueError("Unsupported background method passed. Use background or blank.")

    new_image[
        destination_vertical_lower:destination_vertical_upper,
        destination_horizontal_lower:destination_horizontal_upper
    ] = image[
        source_vertical_lower:source_vertical_upper,
        source_horizontal_lower:source_horizontal_upper
    ]
    return new_image

max_stars_repo_name: csachs/molyso

max_issues_repo_name: csachs/molyso

max_forks_repo_name: csachs/molyso


text:
###########
## Setup ##
###########

import tensorflow as tf
from tensorflow import keras

import numpy as np
import matplotlib.pyplot as plt

print(tf.__version__) # >= 1.12.0

######################
## Download Dataset ##
######################

NUM_WORDS = 10000

(train_data, train_labels), (test_data, test_labels) = keras.datasets.imdb.load_data(num_words=NUM_WORDS)

def multi_hot_sequences(sequences, dimension):
    # Create an all-zero matrix of shape (len(sequences), dimension)
    results = np.zeros((len(sequences), dimension)) # Why does this double enclose? Python quirk?
    for i, word_indices in enumerate(sequences):
        results[i, word_indices] = 1.0 # set specific indices of results[i] to 1s
    return results

train_data = multi_hot_sequences(train_data, dimension=NUM_WORDS)
test_data = multi_hot_sequences(test_data, dimension=NUM_WORDS)

# View multi-hot vectors
#plt.plot(train_data[0])
#plt.show()

#############################
## Demonstrate Overfitting ##
#############################

# Create baseline model

baseline_model = keras.Sequential([
    # `input_shape` is only required here so that `.summary` works.
    keras.layers.Dense(16, activation=tf.nn.relu, input_shape=(NUM_WORDS,)),
    keras.layers.Dense(16, activation=tf.nn.relu),
    keras.layers.Dense(1, activation=tf.nn.sigmoid)
])

baseline_model.compile(
        optimizer='adam',
        loss='binary_crossentropy',
        metrics=['accuracy', 'binary_crossentropy'])

print(baseline_model.summary())

baseline_history = baseline_model.fit(
        train_data,
        train_labels,
        epochs=20,
        batch_size=512,
        validation_data=(test_data, test_labels),
        verbose=2)

# Create a smaller model

smaller_model = keras.Sequential([
    keras.layers.Dense(4, activation=tf.nn.relu, input_shape=(NUM_WORDS,)),
    keras.layers.Dense(4, activation=tf.nn.relu),
    keras.layers.Dense(1, activation=tf.nn.sigmoid)
])

smaller_model.compile(
        optimizer='adam',
        loss='binary_crossentropy',
        metrics=['accuracy', 'binary_crossentropy'])

print(smaller_model.summary())

smaller_history = smaller_model.fit(
        train_data,
        train_labels,
        epochs=20,
        batch_size=512,
        validation_data=(test_data, test_labels),
        verbose=2)

# Create a bigger model

bigger_model = keras.Sequential([
    keras.layers.Dense(256, activation=tf.nn.relu, input_shape=(NUM_WORDS,)),
    keras.layers.Dense(256, activation=tf.nn.relu), # Using 256 to be easier on memory
    keras.layers.Dense(1, activation=tf.nn.sigmoid)
])

bigger_model.compile(
        optimizer='adam',
        loss='binary_crossentropy',
        metrics=['accuracy', 'binary_crossentropy'])

print(bigger_model.summary())

# This kills the PC.
bigger_history = bigger_model.fit(
        train_data,
        train_labels,
        epochs=20,
        batch_size=512,
        validation_data=(test_data, test_labels),
        verbose=2)

#######################################
## Plot training and validation loss ##
#######################################

def plot_history(histories, key='binary_crossentropy'):
    plt.figure(figsize=(16,10))

    for name, history in histories:
        val = plt.plot(
                history.epoch,
                history.history['val_'+key],
                '--',
                label=name.title()+' Val')
        plt.plot(
                history.epoch,
                history.history[key],
                color=val[0].get_color(),
                label=name.title()+' Train')

    plt.xlabel('Epochs')
    plt.ylabel(key.replace('_',' ').title())
    plt.legend()

    plt.xlim([0, max(history.epoch)])
    plt.show()

plot_history([
    ('baseline', baseline_history),
    ('smaller', smaller_history),
    ('bigger', bigger_history)])

################
## Strategies ##
################

# Add weight regularization

l2_model = keras.models.Sequential([
    keras.layers.Dense(16, kernel_regularizer=keras.regularizers.l2(0.001), activation=tf.nn.relu, input_shape=(NUM_WORDS,)),
    keras.layers.Dense(16, kernel_regularizer=keras.regularizers.l2(0.001), activation=tf.nn.relu, input_shape=(NUM_WORDS,)),
    keras.layers.Dense(1, activation=tf.nn.sigmoid)
    ])

l2_model.compile(
        optimizer='adam',
        loss='binary_crossentropy',
        metrics=['accuracy', 'binary_crossentropy'])

l2_model_history = l2_model.fit(
        train_data,
        train_labels,
        epochs=20,
        batch_size=512,
        validation_data=(test_data, test_labels),
        verbose=2)

plot_history([
    ('baseline', baseline_history),
    ('l2', l2_model_history)])

# Add dropout

dropout_model = keras.models.Sequential([
    keras.layers.Dense(16, activation=tf.nn.relu, input_shape=(NUM_WORDS,)),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(16, activation=tf.nn.relu),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(1, activation=tf.nn.sigmoid)
])

dropout_model.compile(
        optimizer='adam',
        loss='binary_crossentropy',
        metrics=['accuracy', 'binary_crossentropy'])

dropout_model_history = dropout_model.fit(
        train_data,
        train_labels,
        epochs=20,
        batch_size=512,
        validation_data=(test_data, test_labels),
        verbose=2)

plot_history([
    ('baseline', baseline_history),
    ('dropout', dropout_model_history)])

# View all three strategies together

plot_history([
    ('baseline', baseline_history),
    ('l2', l2_model_history),
    ('dropout', dropout_model_history)
])

max_stars_repo_name: athix/tensorflow-hello-world

max_issues_repo_name: athix/tensorflow-hello-world

max_forks_repo_name: athix/tensorflow-hello-world


text:
import keras
from keras import backend as K
from keras.models import Sequential
from keras.layers import Conv2D, BatchNormalization, ZeroPadding2D, Lambda, Layer
import pickle
import numpy as np
import cv2

class LRN(Layer):

    def __init__(self, alpha=256,k=0,beta=0.5,n=256, **kwargs):
        super(LRN, self).__init__(**kwargs)
        self.alpha = alpha
        self.k = k
        self.beta = beta
        self.n = n

    def call(self, x, mask=None):

            s = K.shape(x)
            b = s[0]
            r = s[1]
            c = s[2]
            ch = s[3]

            half_n = self.n // 2 # half the local region

            input_sqr = K.square(x) # square the input

            extra_channels = K.zeros((b, r, c, ch + 2 * half_n))
            input_sqr = K.concatenate([extra_channels[:, :, :, :half_n],input_sqr, extra_channels[:, :, :, half_n + ch:]], axis = 3)

            scale = self.k # offset for the scale
            norm_alpha = self.alpha / self.n # normalized alpha
            for i in range(self.n):
                scale += norm_alpha * input_sqr[:, :, :, i:i+ch]
            scale = scale ** self.beta
            x = x / scale

            return x

    def get_config(self):
        config = {"alpha": self.alpha,
                  "k": self.k,
                  "beta": self.beta,
                  "n": self.n}
        base_config = super(LRN, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

def build_cnn(weights):

    model = Sequential()

    model.add(ZeroPadding2D(1, input_shape=(32,32, 1)))
    model.add(Conv2D(32, kernel_size=(3, 3)))
    model.add(BatchNormalization(epsilon=0.0001, scale=False, center=False))
    model.add(Lambda(K.relu))

    model.add(ZeroPadding2D(1))
    model.add(Conv2D(32, kernel_size=(3, 3)))
    model.add(BatchNormalization(epsilon=0.0001, scale=False, center=False))
    model.add(Lambda(K.relu))

    model.add(ZeroPadding2D(1))
    model.add(Conv2D(64, kernel_size=(3, 3), strides=2))
    model.add(BatchNormalization(epsilon=0.0001, scale=False, center=False))
    model.add(Lambda(K.relu))

    model.add(ZeroPadding2D(1))
    model.add(Conv2D(64, kernel_size=(3, 3)))
    model.add(BatchNormalization(epsilon=0.0001, scale=False, center=False))
    model.add(Lambda(K.relu))

    model.add(ZeroPadding2D(1))
    model.add(Conv2D(128, kernel_size=(3, 3), strides=2))
    model.add(BatchNormalization(epsilon=0.0001, scale=False, center=False))
    model.add(Lambda(K.relu))

    model.add(ZeroPadding2D(1))
    model.add(Conv2D(128, kernel_size=(3, 3)))
    model.add(BatchNormalization(epsilon=0.0001, scale=False, center=False))
    model.add(Lambda(K.relu))

    model.add(Conv2D(128, kernel_size=(8, 8)))
    model.add(BatchNormalization(epsilon=0.0001, scale=False, center=False))
    model.add(LRN(alpha=256,k=0,beta=0.5,n=256))

    model.set_weights(weights)

    return model

def build_L2_net(net_name):

    python_net_data = pickle.load(open(net_name + ".pickle", "rb"))
    return build_cnn(python_net_data['weights']), build_cnn(python_net_data['weights_cen']), python_net_data['pix_mean'], python_net_data['pix_mean_cen']

def cal_L2Net_des(net_name, testPatchs, flagCS = False):

    """
    Get descriptors for one or more patches

    Parameters
    ----------
    net_name : string
        One of "L2Net-HP", "L2Net-HP+", "L2Net-LIB", "L2Net-LIB+", "L2Net-ND", "L2Net-ND+", "L2Net-YOS", "L2Net-YOS+",
    testPatchs : array
        A numpy array of image data with deimensions (?, 32, 32, 1), or if using central-surround with deimensions (?, 64, 64, 1)
    flagCS : boolean
        If True, use central-surround network

    Returns
    -------
    descriptor
        Numpy array with size (?, 128) or if using central-surround (?, 256)

    """

    model, model_cen, pix_mean, pix_mean_cen = build_L2_net(net_name)

    # print(model.summary())
    # print(model_cen.summary())

    if flagCS:

        testPatchsCen = testPatchs[:,16:48,16:48,:]
        testPatchsCen = testPatchsCen - pix_mean_cen
        testPatchsCen = np.array([(testPatchsCen[i] - np.mean(testPatchsCen[i]))/(np.std(testPatchsCen[i]) + 1e-12) for i in range(0, testPatchsCen.shape[0])])

        testPatchs = np.array([cv2.resize(testPatchs[i], (32,32), interpolation = cv2.INTER_CUBIC) for i in range(0, testPatchs.shape[0])])
        testPatchs = np.expand_dims(testPatchs, axis=-1)

    testPatchs = testPatchs - pix_mean
    testPatchs = np.array([(testPatchs[i] - np.mean(testPatchs[i]))/(np.std(testPatchs[i]) + 1e-12) for i in range(0, testPatchs.shape[0])])

    res = np.reshape(model.predict(testPatchs), (testPatchs.shape[0], 128))

    if flagCS:

        resCen = np.reshape(model_cen.predict(testPatchsCen), (testPatchs.shape[0], 128))

        return np.concatenate((res, resCen), 1)

    else:

        return res

class L2Net:

    def __init__(self, net_name, flagCS = False):
        model, model_cen, pix_mean, pix_mean_cen = build_L2_net(net_name)
        self.flagCS = flagCS
        self.model = model
        self.model_cen = model_cen
        self.pix_mean = pix_mean
        self.pix_mean_cen = pix_mean_cen

    def calc_descriptors(self, patches):
        if self.flagCS:

            patchesCen = patches[:,16:48,16:48,:]
            patchesCen = patchesCen - self.pix_mean_cen
            patchesCen = np.array([(patchesCen[i] - np.mean(patchesCen[i]))/(np.std(patchesCen[i]) + 1e-12) for i in range(0, patchesCen.shape[0])])

            patches = np.array([cv2.resize(patches[i], (32,32), interpolation = cv2.INTER_CUBIC) for i in range(0, patches.shape[0])])
            patches = np.expand_dims(patches, axis=-1)

        patches = patches - self.pix_mean
        patches = np.array([(patches[i] - np.mean(patches[i]))/(np.std(patches[i]) + 1e-12) for i in range(0, patches.shape[0])])

        res = np.reshape(self.model.predict(patches), (patches.shape[0], 128))

        if self.flagCS:

            resCen = np.reshape(self.model_cen.predict(patchesCen), (patches.shape[0], 128))

            return np.concatenate((res, resCen), 1)

        else:

            return res

# data = np.full((1,64,64,1), 0.)

# result = cal_L2Net_des("L2Net-HP+", data, flagCS=True)

# print(result)

max_stars_repo_name: virtualgraham/fgraph

max_issues_repo_name: virtualgraham/fgraph

max_forks_repo_name: virtualgraham/fgraph


text:
#!/usr/bin/env python
"""
This script tests the python class interface
"""

from __future__ import absolute_import, division, print_function

# standard imports:
import os
import sys

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

sys.path.append(os.path.join(os.path.dirname(__file__), "../../.."))
import pyshtools

pyshtools.utils.figstyle()

# ==== MAIN FUNCTION ====

def main():
    example1()
    example2()

# ==== EXAMPLES ====
def example1():
    # generate cap window
    lmax = 20
    nwin = 20
    theta = 25.
    cap = pyshtools.SHWindow.from_cap(theta, lmax, nwin=nwin)
    cap.info()
    cap.plot_windows(20, show=False, fname='cap_tapers.png')
    cap.plot_coupling_matrix(30, nwin=5, show=False, fname='cap_coupling.png')

# ==== EXAMPLES ====
def example2():
    # generate cap window
    lmax = 15
    nwins = 15

    coeffs = pyshtools.SHCoeffs.from_file('../../ExampleDataFiles/srtmp300.msl')
    topo = coeffs.expand(grid='DH2')
    dh_mask = topo.data > 0.
    print(dh_mask.shape)
    region = pyshtools.SHWindow.from_mask(dh_mask, lmax, nwins)
    region.info()
    region.plot_windows(nwins, show=False, fname='continent_tapers.png')
    region.plot_coupling_matrix(30, 5, show=False,
                                fname='continent_coupling.png')

# ==== EXECUTE SCRIPT ====
if __name__ == "__main__":
    main()


text:
"""
"""
from copy import deepcopy
from typing import Union, Optional, Sequence, Tuple, NoReturn

import numpy as np
import pandas as pd
import torch
from torch import nn
from torch import Tensor

try:
    import torch_ecg
except ModuleNotFoundError:
    import sys
    from os.path import dirname, abspath
    sys.path.insert(0, dirname(dirname(dirname(abspath(__file__)))))

from torch_ecg.cfg import CFG
from torch_ecg.models.ecg_crnn import ECG_CRNN

from cfg import ModelCfg

__all__ = [
    "ECG_CRNN_CINC2021",
]

class ECG_CRNN_CINC2021(ECG_CRNN):
    """
    """
    __DEBUG__ = False
    __name__ = "ECG_CRNN_CINC2021"

    def __init__(self, classes:Sequence[str], n_leads:int, config:Optional[CFG]=None) -> NoReturn:
        """ finished, checked,

        Parameters
        ----------
        classes: list,
            list of the classes for classification
        n_leads: int,
            number of leads (number of input channels)
        config: dict, optional,
            other hyper-parameters, including kernel sizes, etc.
            ref. the corresponding config file
        """
        model_config = CFG(deepcopy(ModelCfg))
        model_config.update(deepcopy(config) or {})
        super().__init__(classes, n_leads, model_config)

    @torch.no_grad()
    def inference(self,
                  input:Union[np.ndarray,Tensor],
                  class_names:bool=False,
                  bin_pred_thr:float=0.5) -> Tuple[Union[np.ndarray,pd.DataFrame],np.ndarray]:
        """ finished, checked,

        auxiliary function to `forward`, for CINC2021,

        Parameters
        ----------
        input: ndarray or Tensor,
            input tensor, of shape (batch_size, channels, seq_len)
        class_names: bool, default False,
            if True, the returned scalar predictions will be a `DataFrame`,
            with class names for each scalar prediction
        bin_pred_thr: float, default 0.5,
            the threshold for making binary predictions from scalar predictions

        Returns
        -------
        pred: ndarray or DataFrame,
            scalar predictions, (and binary predictions if `class_names` is True)
        bin_pred: ndarray,
            the array (with values 0, 1 for each class) of binary prediction

        NOTE that when `input` is ndarray, one should make sure that it is transformed,
        e.g. bandpass filtered, normalized, etc.
        """
        if "NSR" in self.classes:
            nsr_cid = self.classes.index("NSR")
        elif "426783006" in self.classes:
            nsr_cid = self.classes.index("426783006")
        else:
            nsr_cid = None
        _device = next(self.parameters()).device
        _dtype = next(self.parameters()).dtype
        _input = torch.as_tensor(input, dtype=_dtype, device=_device)
        if _input.ndim == 2:
            _input = _input.unsqueeze(0)  # add a batch dimension
        # batch_size, channels, seq_len = _input.shape
        pred = self.forward(_input)
        pred = self.sigmoid(pred)
        bin_pred = (pred>=bin_pred_thr).int()
        pred = pred.cpu().detach().numpy()
        bin_pred = bin_pred.cpu().detach().numpy()
        for row_idx, row in enumerate(bin_pred):
            row_max_prob = pred[row_idx,...].max()
            if row_max_prob < ModelCfg.bin_pred_nsr_thr and nsr_cid is not None:
                bin_pred[row_idx, nsr_cid] = 1
            elif row.sum() == 0:
                bin_pred[row_idx,...] = \
                    (((pred[row_idx,...]+ModelCfg.bin_pred_look_again_tol) >= row_max_prob) & (pred[row_idx,...] >= ModelCfg.bin_pred_nsr_thr)).astype(int)
        if class_names:
            pred = pd.DataFrame(pred)
            pred.columns = self.classes
            # pred["bin_pred"] = pred.apply(
            #     lambda row: np.array(self.classes)[np.where(row.values>=bin_pred_thr)[0]],
            #     axis=1
            # )
            pred["bin_pred"] = ""
            for row_idx in range(len(pred)):
                pred.at[row_idx, "bin_pred"] = \
                    np.array(self.classes)[np.where(bin_pred[row_idx]==1)[0]].tolist()
        return pred, bin_pred

    @torch.no_grad()
    def inference_CINC2021(self,
                           input:Union[np.ndarray,Tensor],
                           class_names:bool=False,
                           bin_pred_thr:float=0.5) -> Tuple[Union[np.ndarray,pd.DataFrame],np.ndarray]:
        """
        alias for `self.inference`
        """
        return self.inference(input, class_names, bin_pred_thr)