text:
from .database import kirkland_params
from .optics import *
from .utils import *
from .cuda_kernels import ProbeKernels, PotentialKernels
import numpy as np
from scipy.special import k0
import multiprocessing as mp
import ctypes
import sys
from warnings import warn, catch_warnings, simplefilter
import os

try:
    import pyfftw
except:
    pyfftw=None
import pycuda
import pycuda.driver as cuda
import skcuda.fft as skfft
import skcuda.cufft as cufft
from time import time
import h5py
from mpi4py import MPI

XYZ_dtype = [('atomic_number', 'i'), ('x', 'f'), ('y', 'f'), ('z', 'f'), ('occ', 'f'), ('DW', 'f')]

def catch_warn():
    warn('complex_casting', RuntimeWarning)

def unwrap(args):
    (msa, params), kwargs = args
    method_name = kwargs.pop('method')
    if method_name == 'build_slices':
        return msa.make_slice(params)
    elif method_name == 'propagate_beams':
        return msa.propagate_beam(params)
    elif method_name == 'build_probes':
        return msa.build_probe(probe_position=params[0], probe_dict=params[1])

class MSA(object):
    def __init__(self, energy, semi_angle, supercell, sampling=np.array([512, 512]), max_angle=None, verbose=False,
                 debug=False, output_dir='', debye_waller=True):
        self.E = energy
        self.Lambda = voltage2Lambda(self.E*1e3)
        self.semi_ang = semi_angle
        self.verbose = verbose
        self.debug = debug
        self.output_dir = output_dir
        self.debye_waller = debye_waller
        # Load and set supercell properties
        if isinstance(supercell, np.ndarray):
            try:
                check = [supercell[0][field] for field in ['x', 'y', 'z', 'occ', 'DW', 'atomic_number']]
                supercell_arr = np.copy(supercell)
            except ValueError as err:
                warn('Supercell numpy array does not have required fields.')
                raise err
        elif isinstance(supercell, str):
            if os.path.exists(supercell):
                try:
                    supercell_arr = np.genfromtxt(supercell, skip_footer=1, skip_header=2, dtype=XYZ_dtype)
                except IOError as err:
                    warn('Unable to load supercell file into numpy array.')
                    raise err
            else:
                warn('Path %s was not found.' % supercell)
                sys.exit(0)
        else:
            warn('supercell format was not recognized.')
            sys.exit(0)
        self.supercell_xyz = np.column_stack([supercell_arr['x'], supercell_arr['y'], supercell_arr['z']]).astype(np.float32)
        self.supercell_Z = supercell_arr['atomic_number']
        self.supercell_dw = supercell_arr['DW']
        self.supercell_occ = supercell_arr['occ']

        # Adding random displacements due to thermal effects
        unique_Z = np.unique(self.supercell_Z)
        for Z in unique_Z:
            indx = np.where(self.supercell_Z == Z)[0]
            dw = self.supercell_dw[indx][0]
            displacements = float(self.debye_waller) * dw * np.random.standard_normal(size=self.supercell_xyz[indx].shape)
            self.supercell_xyz[indx] += displacements

        # Set simulation parameters
        self.dims = self.supercell_xyz.max(0) - self.supercell_xyz.min(0)
        if max_angle is None:
            self.sampling = sampling
            self.kmax = np.min(self.sampling/self.dims[:2])
            self.max_ang = self.kmax * self.Lambda
        else:
            self.max_ang = max_angle
            self.kmax = self.max_ang / self.Lambda
            self.sampling = np.floor(self.kmax * self.dims[:2]).astype(np.int)
            self.sampling = self.sampling[::-1]
        self.pix_size = self.dims[:2][::-1] / self.sampling
        self.kpix_size = self.kmax/self.sampling
        self.sigma = sigma_int(self.E*1e3)
        self.print_verbose('Simulation Parameters:\nSupercell dimensions xyz:%s (Å)\nReal, Reciprocal space pixel sizes:%s Å, %s 1/Å'
              '\nMax angle: %2.2f (rad)\nSampling in real and reciprocal space: %s pixels,\nThermal Effects: %s' %
              (format(np.round(self.dims, 2)), format(np.round(self.pix_size, 2)), format(np.round(self.kpix_size, 2)),
               self.max_ang, format(self.sampling), format(self.debye_waller)))

    def print_debug(self, *args, **kwargs):
        if self.debug:
            print(*args, **kwargs)

    def print_verbose(self, *args, **kwargs):
        if self.verbose:
            print(*args, **kwargs)

    def calc_atomic_potentials(self, potential_range=8, oversample=2, kirkland=True):
        if kirkland:
            self.scattering_params = kirkland_params
            #TODO: Figure out how to calculate scattering potential using different methods
        start = potential_range / 2
        step_x, step_y = 1.j * np.floor(potential_range/self.pix_size).astype(np.int) * oversample
        grid_x, grid_y = np.mgrid[-start:start:step_x, -start:start:step_y]
        self.cached_pots = dict()
        for Z in np.unique(self.supercell_Z):
            self.cached_pots[Z] = self._get_potential(grid_x, grid_y, self.scattering_params[Z - 1], oversample)

    def build_potential_slices(self, slice_thickness):
        self.slice_t = slice_thickness
        num_slices = np.int(np.floor(self.dims[-1] / slice_thickness))
        tasks = [((self, slice_num), {'method': 'build_slices'}) for slice_num in range(num_slices)]
        processes = min(mp.cpu_count(), num_slices)
        chunk = np.int(np.floor(num_slices / processes))
        pool = mp.Pool(processes=processes, maxtasksperchild=1)
        jobs = pool.imap(unwrap, tasks, chunksize=chunk)
        potential_slices = np.array([j for j in jobs])
        pool.close()
        self.potential_slices = potential_slices
        self.print_verbose('Built %d potential slices with shape:%s pixels' % (self.potential_slices.shape[0],
                                                                  format(self.potential_slices.shape[1:])))

    def make_slice(self, args):
        slice_num = args
        mask = np.logical_and(self.supercell_xyz[:, -1] >= slice_num * self.slice_t,
                              self.supercell_xyz[:, -1] < (slice_num + 1) * self.slice_t)
        arr_slice = np.zeros(self.sampling, dtype=np.float32)
        for Z, xyz in zip(self.supercell_Z[mask], self.supercell_xyz[mask]):
            pot = self.cached_pots[Z]
            y_pix, x_pix = xyz[:2][::-1] / self.pix_size
            y_start, y_end = int(y_pix - pot.shape[0] / 2), int(y_pix + pot.shape[0] / 2)
            x_start, x_end = int(x_pix - pot.shape[1] / 2), int(x_pix + pot.shape[1] / 2)
            repl_y = slice(max(y_start, 0), y_end)
            repl_x = slice(max(x_start, 0), x_end)
            offset_y = abs(min(y_start, 0))
            offset_x = abs(min(x_start, 0))
            repl_shape = arr_slice[repl_y, repl_x].shape
            arr_slice[repl_y, repl_x] += pot[offset_y:repl_shape[0] + offset_y, offset_x:repl_shape[1] + offset_x]
        return arr_slice

    @staticmethod
    def _get_potential(grid_x, grid_y, params, oversample):
        eps = 0.1  # Å cutoff for Bessel function
        a_0 = 0.529  # Å Bohr radius
        e = 14.4  # V.Å electron charge
        r = np.sqrt(grid_x ** 2 + grid_y ** 2)
        prefac_1 = 4 * np.pi ** 2 * a_0 * e
        coeff = params
        a_coeff, b_coeff, c_coeff, d_coeff = coeff
        bessel_arg = 2 * np.pi * r[:, :, np.newaxis] * np.sqrt(b_coeff)
        bessel_eval = k0(bessel_arg + eps)
        sum_v = prefac_1 * a_coeff * bessel_eval + prefac_1 / 2. * c_coeff / d_coeff * np.exp(-np.pi ** 2 *
                                                                                r[:, :, np.newaxis] ** 2 / d_coeff)
        v = sum_v.sum(-1)
        potential = bin_2d_array(v, dwnspl=oversample, mode='mean')
        potential -= potential.min()
        return potential.astype(np.float32)

    def build_probe(self, probe_position=np.array([0., 0.]), probe_dict={'smooth_apert': True, 'apert_smooth': 50, 'spherical_phase': True, 'aberration_dict':
                        {'C1': 0., 'C3': 0., 'C5': 0.},'scherzer': True}
                    ):
        self.probe_dict = probe_dict
        k_y, k_x = np.mgrid[-self.kmax/2: self.kmax/2: 1j*self.sampling[0],
                            -self.kmax/2: self.kmax/2: 1j * self.sampling[1]]
        k_rad = np.sqrt(k_x ** 2 + k_y ** 2)
        k_semi = self.semi_ang/self.Lambda

        # aperture function
        if probe_dict['smooth_apert']:
            aperture = 1 / (1 + np.exp(-2 * probe_dict['apert_smooth'] * (k_semi - k_rad)))
        else:
            aperture = np.heaviside(k_semi - k_rad, 0.5)

        # aberration
        if probe_dict['spherical_phase']:
            phase_error = spherical_phase_error(k_rad, self.Lambda, probe_dict['scherzer'],
                                                **probe_dict['aberration_dict'])

        else:
            pass
            #TODO: implement non-rotationally invariant phase error

        # probe wavefunction
        psi_k = aperture * phase_error
        psi_k = psi_k.astype(np.complex64)
        y, x = probe_position
        kr = k_x * x + k_y * y
        phase_shift = np.exp(2 * np.pi * 1.j * kr).astype(np.complex64)
        psi_x = pyfftw.interfaces.numpy_fft.ifft2(psi_k * phase_shift)
        psi_x = pyfftw.interfaces.numpy_fft.fftshift(psi_x)
        # TODO: need to make fft library choice optional
        # psi_x = np.fft.ifft2(psi_k * phase_shift, norm='ortho')
        # psi_x = np.fft.fftshift(psi_x)
        psi_x /= np.sqrt(np.sum(np.abs(psi_x) ** 2))
        return psi_x.astype(np.complex64), psi_k.astype(np.complex64), aperture.astype(np.float32)

    def build_propagator(self):
        k_y, k_x = np.mgrid[-self.kmax / 2: self.kmax / 2: 1j * self.sampling[0],
                   -self.kmax / 2: self.kmax / 2: 1j * self.sampling[1]]
        k_rad_sq = k_x ** 2 + k_y ** 2
        propag = np.exp(-np.pi * 1.j * self.Lambda * self.slice_t * k_rad_sq)
        return propag.astype(np.complex64)

    @staticmethod
    def bandwidth_limit_mask(arr_shape, radius=0.5):
        # assumes square image
        grid_x, grid_y = np.mgrid[-arr_shape[0] // 2:arr_shape[0] // 2, -arr_shape[1] // 2:arr_shape[1] // 2]
        r_grid = np.sqrt(grid_x ** 2 + grid_y ** 2)
        bl_mask = np.heaviside(max(arr_shape[0], arr_shape[1]) * radius - r_grid, 0)
        return bl_mask.astype(np.float32)

    def generate_probe_positions(self, probe_step=np.array([0.1, 0.1]), probe_range=np.array([[0., 1.0], [0., 1.0]]), grid_steps=None):
        if grid_steps is not None:
            grid_range_start = (0.5 - self.dims[:2]/2)/2
            grid_range_stop = (0.5 + self.dims[:2]/2)/2
            x_pos, y_pos = np.mgrid[grid_range_start[0]:grid_range_stop[0]:-1j*grid_steps[0],
                                    grid_range_start[1]:grid_range_stop[1]:-1j*grid_steps[1]]
            grid_steps_x, grid_steps_y = grid_steps.astype(np.int)
            grid_range_x = np.array([grid_range_start[0], grid_range_stop[0]])
            grid_range_y = np.array([grid_range_start[1], grid_range_stop[1]])
            probe_pos = np.array([[-x, y] for y, x in zip(y_pos.flatten()[::-1], x_pos.flatten())])

        else:
            grid_steps_x, grid_steps_y = np.floor(np.diff(probe_range).flatten() * self.dims[:2] / probe_step).astype(np.int)
            grid_range_x, grid_range_y = [(probe_range[i] - np.ones((2,)) * 0.5) * self.dims[i]
                                        for i in range(2)]
            y_pos, x_pos = np.mgrid[grid_range_y[0]: grid_range_y[1]: -1j * grid_steps_y,
                        grid_range_x[0]: grid_range_x[1]: -1j * grid_steps_x]
            probe_pos = np.array([[y, -x] for y, x in zip(y_pos.flatten()[::-1], x_pos.flatten())])
        self.grid_steps = np.array([grid_steps_x, grid_steps_y])
        self.grid_range = np.array([grid_range_x, grid_range_y]).flatten()
        self.probe_positions = probe_pos
        self.num_probes = np.int32(probe_pos.shape[0])

    def multislice(self, probe_pos=np.array([0., 0.]), probe_grid=True, save_probes=True, bandwidth=1 / 3):
        # check for slices
        if isinstance(self.potential_slices, np.ndarray) is False:
            warn('Potential slices must be calculated first before calling multi_slice')
            return
        # check for probe
        if isinstance(self.probe_dict, dict) is False:
            warn('Probe wave function must be initialized first before calling multi_slice')
            return
        if probe_grid:
            # Define heuristics for python multiprocessing and FFTW multi-threading
            self.fftw_threads = int(self.sampling.max() // 512)
            processes = min(mp.cpu_count(), self.probe_positions.shape[0]) // self.fftw_threads
            chunk = int(np.floor(self.probe_positions.shape[0] / processes))

            # Put the potential slices in shared memory so all workers access it (asynchronously)
            global shared_slices
            shared_slices = mp.Array(ctypes.c_float, self.potential_slices.size, lock=False)
            temp = np.frombuffer(shared_slices, dtype=np.float32)
            for (i, pot) in enumerate(self.potential_slices):
                temp[i * pot.size:(i + 1) * pot.size] = pot.flatten().astype(np.float32)

            tasks = (((self, (probe_num, pos, save_probes, probe_grid, bandwidth)), {'method': 'propagate_beams'})
                     for (probe_num, pos) in enumerate(self.probe_positions))
            pool = mp.Pool(processes=processes, initargs=(shared_slices,))
            jobs = pool.map(unwrap, tasks, chunksize=chunk)
            trans_probes = np.array([j for j in jobs])
            pool.close()
            pool.join()
        else:
            trans_probes = self.propagate_beam([None, probe_pos, save_probes, probe_grid, bandwidth])

        self.print_verbose('Propagated %d probe wavefunctions' % trans_probes.shape[0])
        self.trans_probes = trans_probes
        return self.trans_probes

    def propagate_beam(self, args):
        probe_num, probe_pos, save_probes, probe_grid, bandwidth = args
        self.print_debug('received params.')
        propag = self.build_propagator()
        blim_mask = self.bandwidth_limit_mask(propag.shape, radius=bandwidth)
        probe, _, _ = self.build_probe(probe_pos, self.probe_dict)
        probes = []
        probe_last = probe
        if probe_grid:
            slices = np.frombuffer(shared_slices, dtype=np.float32)
            slices = slices.reshape(self.potential_slices.shape)
            self.print_debug('fetched data from shared memory.')
        else:
            slices = self.potential_slices
        slices = np.exp(1.j * self.sigma * slices).astype(np.complex64)

        for (i, trans) in enumerate(slices[::-1]):
            t_psi = pyfftw.byte_align(trans * probe_last)
            fft_fwd = pyfftw.builders.fft2(t_psi, threads=self.fftw_threads, avoid_copy=True)
            temp = pyfftw.byte_align(fft_fwd() * blim_mask * propag)
            fft_bwd = pyfftw.builders.ifft2(temp, threads=self.fftw_threads, avoid_copy=True)
            probe_last = fft_bwd()
            if save_probes:
                probes.append(probe_last)
        self.print_debug('finished beam propagation.')

        if save_probes:
            np.save(os.path.join(self.output_dir, 'probes_%d.npy') % probe_num, np.array(probes), allow_pickle=False)
            del probes
            self.print_verbose('finished with probe position: %s' % format(probe_pos))

        return probe_last

    def check_simulation(self):
        prob = np.sum([np.abs(probe)**2 for probe in self.trans_probes], axis=(1, 2))
        max_val = prob.max()
        min_val = prob.min()
        print('Max (Min) Integrated Intensity: %2.2f (%2.2f)' % (max_val, min_val))
        if max_val > 1.0 or min_val < 0.95:
            print('Significant deviation of the probability from unity is found.\n'
                  'Change the sampling and/or slice thickness!')
        return prob

class MSAHybrid(MSA):
    def setup_device(self, gpu_rank=0):
        global ctx
        cuda.init()
        dev = cuda.Device(gpu_rank)
        ctx = dev.make_context()
        # ctx.attach()
        self.gpu_rank = gpu_rank

        import atexit
        def _clean_up():
            global ctx
            if ctx is not None:
                try:#global ctx
                    #ctx.push()
                    ctx.pop()
                    ctx.detach()
                    #ctx = None
                except:
                    pass
            from pycuda.tools import clear_context_caches
            clear_context_caches()

        atexit.register(_clean_up)
        return ctx # return context in case of manual clean-up

    def plan_simulation(self, num_probes=None):
        if num_probes is None:
            num_probes = self.num_probes
        self.print_verbose('Simulation requested %d probes simultaneously.' % self.num_probes)
        free_mem, tot_mem = pycuda.driver.mem_get_info()
        free_mem = free_mem/1024e6  # in GB
        mem_alloc = num_probes * np.prod(self.sampling) * 8 / 1024e6 + self.potential_slices.nbytes / 1024e6
        self.print_verbose('mem_alloc: %2.3f' % mem_alloc)
        self.print_verbose('free_mem: %2.3f' % free_mem)
        self.max_probes = num_probes
        mod = False
        while mem_alloc > 0.5 * free_mem:
            mod = True
            num_probes = int(num_probes * 0.75)
            mem_alloc = num_probes * np.prod(self.sampling) * 8 / 1024e6 + self.potential_slices.nbytes / 1024e6
            self.max_probes = num_probes
        if mod: self.print_verbose('Device can hold at most %d probes simultaneously.' % self.max_probes)

    def multislice(self, save_probes=True, bandwidth=1/3):
        # not supporting a single probe!!
        # check for potential slices
        if isinstance(self.potential_slices, np.ndarray) is False:
            warn('Potential slices must be calculated first before calling multi_slice')
            return
        # check for probe
        if isinstance(self.probe_dict, dict) is False:
            warn('Probe wave function must be initialized first before calling multi_slice')
            return
        # check for probe positions
        if isinstance(self.probe_positions, np.ndarray) is False:
            warn('probe positions must be initialized first before calling multi_slice')
            return

        # Initialize needed data
        slices = np.exp(1.j * self.sigma * self.potential_slices).astype(np.complex64)
        propag = self.build_propagator()
        mask = self.bandwidth_limit_mask(propag.shape, radius=bandwidth)
        t = time()
        self.probes = self.build_probes_cpu()
        sim_t = time() - t
        self.print_verbose('Spent %2.4f s building %d probes on cpu' % (sim_t, self.max_probes))

        # Copy over to device
        trans_gpu = pycuda.gpuarray.to_gpu_async(slices)
        mask_propag_gpu = pycuda.gpuarray.to_gpu_async(mask * propag)
        probes_gpu = pycuda.gpuarray.to_gpu_async(self.probes)

        # Setup fft plans
        # TODO: tile multiple fft plans
        t = time()
        plan = skfft.Plan(propag.shape, np.complex64, np.complex64, batch=self.max_probes)
        for slice_num in range(trans_gpu.shape[0]):
            for z_slice in range(probes_gpu.shape[0]):
                probes_gpu[z_slice] *= trans_gpu[slice_num]
            cufft.fft(probes_gpu, probes_gpu, plan, True)
            for z_slice in range(probes_gpu.shape[0]):
                probes_gpu[z_slice] *= mask_propag_gpu
            cufft.ifft(probes_gpu, probes_gpu, plan, False)
        cufft.fft(probes_gpu, probes_gpu, plan, True)  #return probe wavefunctions in reciprocal space
        sim_t = time() - t
        self.print_verbose('Propagated %d probes in %2.4f s' % (self.max_probes, sim_t))
        self.probes = probes_gpu.get()

        # free up device memory
        trans_gpu.gpudata.free()
        mask_propag_gpu.gpudata.free()
        probes_gpu.gpudata.free()

        # destroy cufft object
        cufft.cufftDestroy(plan.handle)
        return self.probes

    def build_probes_cpu(self):
        processes = min(mp.cpu_count(), self.probe_positions.shape[0]) // 4
        chunk = int(np.floor(self.probe_positions.shape[0] / processes))
        tasks = (((self, (pos, self.probe_dict)), {'method': 'build_probes'})
                 for pos in self.probe_positions)
        pool = mp.Pool(processes=processes)
        jobs = pool.map(unwrap, tasks, chunksize=chunk)
        probes = np.array([j[0] for j in jobs])
        pool.close()
        pool.join()
        return probes

class MSAGPU(MSAHybrid):

    @staticmethod
    def clean_up(ctx=None, vars=None):
        if vars is not None:
            for var in vars:
                if var is not None:
                    var.free()
        if ctx is not None:
            ctx.pop()
            ctx.detach()
            ctx = None
            from pycuda.tools import clear_context_caches
            clear_context_caches()

#     def build_potential_slices(self, slice_thickness):
# #         # find number of slices and atomic sites per slice
#         self.slice_t = slice_thickness
#         self.num_slices = np.int32(np.floor(self.dims[-1] / slice_thickness))
# #         masks = [np.logical_and(self.supercell_xyz[:, -1] >= slice_num * self.slice_t,
# #                         self.supercell_xyz[:, -1] < (slice_num + 1) * self.slice_t)
# #                         for slice_num in range(self.num_slices)]
# #         # remap Z numbers to indices in cached atomic potential dictionary
# #         unique_Z = np.unique(self.supercell_Z)
# #         supercell_Z_idx = np.array([np.argmax(np.equal(unique_Z, Z_val)) for Z_val in self.supercell_Z])
# #         Z_arr = np.array([supercell_Z_idx[mask] for mask in masks])
# #         supercell_pix = self.supercell_xyz[:,:2]/self.pix_size[::-1]
# #         yx_arr = np.array([supercell_pix[mask] for mask in masks])

# #         # stack the input for all slices into a single 1-d array
# #         Zxy_input = np.array([np.column_stack([Z, yx[:,1], yx[:,0]]).astype(np.int32).flatten()
# #                                             for Z, yx in zip(Z_arr, yx_arr)])

# #         # pad sites array to have equal num of element per slice
# #         shapes = np.array([itm.shape for itm in Zxy_input])
# #         pads = [(0, shapes.max() - itm.shape) for itm in Zxy_input]
# #         padded = [np.pad(itm,pad,'constant',constant_values=-1) for pad, itm in zip(pads, Zxy_input)]
# #         Zxy_input = np.vstack([np.pad(itm,pad,'constant',constant_values=-1) for pad, itm in zip(pads, Zxy_input)])

# #         # stack atomic potentials of unique elements
# #         atom_pot_stack = np.array([self.cached_pots[uq_Z] for uq_Z in unique_Z]).astype(np.float32)

# #         # compile/load cuda kernels
# #         self._load_potential_kernels([atom_pot_stack.shape[1], atom_pot_stack.shape[2]],
# #                                                 Zxy_input.shape[-1])

# #         # allocate memory
# #         atom_pot_stack_d = cuda.to_device(atom_pot_stack)
# #         sites_d = cuda.to_device(Zxy_input)
# #         self.potential_slices = cuda.aligned_empty((int(self.num_slices),
# #                      int(self.sampling[0]), int(self.sampling[1])), np.complex64)
#         self.vars = []
# #         self.vars.append(self.potential_slices.base)
# #         self.potential_slices = cuda.register_host_memory(self.potential_slices)
#         potential_slices_d = cuda.to_device(self.potential_slices)

#         # store device allocation ref for later use
#         self.pot_dev_ptr = potential_slices_d

#         # get block, grid dimensions
# #         block, grid = self._get_blockgrid([self.sampling[1], self.sampling[0], self.num_slices],
# #                     mode='3D')
# #         build_potential = self.pot_kernels['build_potential']

# #         # build potential
# #         build_potential(potential_slices_d, atom_pot_stack_d, sites_d,
# #                         np.float32(self.sigma), block=block, grid=grid)
#         ctx.synchronize()
# #         cuda.memcpy_dtoh_async(self.potential_slices, potential_slices_d, cuda.Stream())

# #         # free gpu memory
# #         sites_d.free()
# #         atom_pot_stack_d.free()
#         self.print_verbose('Built %d potential slices with shape:%s pixels' % (self.potential_slices.shape[0],
#                                                                   format(self.potential_slices.shape[1:])))
#         # unregister host memory
#         #self.potential_slices.base.unregister()
#         #self.potential_slices.base.free()

    def _load_kernels(self):
        try:
            probe_kernels = ProbeKernels(sampling=self.sampling)
            self.kernels = probe_kernels.kernels
            self.print_verbose('CUDA C/C++ Probe Kernels compiled successfully')
        except cuda.CompileError:
            warn('CUDA C/C++ Probe Kernels did not compile successfully')
            raise cuda.CompileError

    def _load_potential_kernels(self, potential_shape, sites_size ):
        try:
            potential_kernels = PotentialKernels(sampling=self.sampling,
            num_slices=self.num_slices, potential_shape=potential_shape,
            sites_size=sites_size)
            self.pot_kernels = potential_kernels.kernels
            self.print_verbose('CUDA C/C++ Potential Kernels compiled successfully')
        except cuda.CompileError:
            warn('CUDA C/C++ Potential Kernels did not compile successfully')
            raise cuda.CompileError

    def build_probe(self, probe_dict={'smooth_apert': True, 'apert_smooth': 50, 'spherical_phase': True,
                                      'aberration_dict': {'C1': 0., 'C3': 0., 'C5': 0.},'scherzer': True}):
        # load kernels and prepare func args
        self._load_kernels()
        aber_dict = probe_dict['aberration_dict']
        k_semi = np.float32(self.semi_ang / self.Lambda)
        k_max, Lambda, C1, C3, C5 = [np.float32(itm) for itm in [self.kmax, self.Lambda, aber_dict['C1'],
                                                                 aber_dict['C3'], aber_dict['C5']]]
        if probe_dict['scherzer']:
            C1, _ = scherzer_params(self.Lambda, aber_dict['C3'])
            C1 = np.float32(C1)

        # grab appropriate kernels
        self.multwise_2d_func = self.kernels['mult_wise_c2d_re2d']
        fftshift_func = self.kernels['fftshift_2d']
        if probe_dict['smooth_apert']:
            apert_func = self.kernels['soft_aperture']
        else:
            apert_func = self.kernels['hard_aperture']
        if probe_dict['spherical_phase']:
            phase_func = self.kernels['spherical_phase_error']
        else:
            pass
            # TODO: implement non-rotationally invariant phase error kernel

        # define block/grid threads
        shape_x = np.int32(self.sampling[1])
        shape_y = np.int32(self.sampling[0])
        block, grid = self._get_blockgrid([shape_x, shape_y], mode='2D')

        # allocate memory
        self.apert = np.empty(self.sampling, dtype=np.float32)
        #self.apert = cuda.register_host_memory(self.apert)
        apert_d = cuda.mem_alloc(self.apert.nbytes)
        self.psi_k = np.empty(self.sampling, dtype=np.complex64)
        #self.psi_k = cuda.register_host_memory(self.psi_k)
        self.psi = np.empty_like(self.psi_k)
        #self.psi = cuda.register_host_memory(self.psi)
        psi_k_d = cuda.mem_alloc(self.psi_k.nbytes)
        psi_x_d = cuda.mem_alloc(self.psi_k.nbytes)

        # build a probe in k-space
        apert_func(apert_d, k_max, k_semi, shape_x, shape_y, block=block, grid=grid, shared=0, stream=cuda.Stream())
        #ctx.synchronize()
        phase_func(psi_k_d, k_max, Lambda, C1, C3, C5, shape_x, shape_y, block=block, grid=grid, shared=0, stream=cuda.Stream())
        #ctx.synchronize()
        self.multwise_2d_func(psi_k_d, apert_d, shape_x, shape_y, block=block, grid=grid)
        #ctx.synchronize()
        cuda.memcpy_dtoh_async(self.psi_k, psi_k_d, cuda.Stream())
        cuda.memcpy_dtoh_async(self.apert, apert_d, cuda.Stream())

        # build probe in x-space
        fft_plan = skfft.Plan(self.psi_k.shape, np.complex64, np.complex64, batch=1)
        cufft.cufftExecC2C(fft_plan.handle, int(psi_k_d), int(psi_x_d), cufft.CUFFT_INVERSE)
        ctx.synchronize()
        fftshift_func(psi_x_d, shape_y, block=block, grid=grid)
        #ctx.synchronize()
        cuda.memcpy_dtoh_async(self.psi, psi_x_d, cuda.Stream())
        self.normalization = np.sum(np.abs(self.psi, dtype=np.float64)**2)
        self.psi /= np.sqrt(self.normalization)
        cuda.memcpy_dtoh_async(self.psi_k, psi_k_d, cuda.Stream())

        # free gpu mem
        psi_x_d.free()
        psi_k_d.free()
        apert_d.free()
        cufft.cufftDestroy(fft_plan.handle)

    @staticmethod
    def _get_blockgrid(shapes, mode='2D'):
        # define block/grid threads
        if mode == '3D':
            shape_x = shapes[0]
            shape_y = shapes[1]
            shape_z = shapes[2]
            blk_zsize = 1
            blk_xsize = 32
            blk_ysize = 32
            grid_xsize = int(round((shape_x + blk_xsize - 1) / blk_xsize))
            grid_ysize = int(round((shape_y + blk_ysize - 1) / blk_ysize))
            grid_zsize = int(round((shape_z + blk_zsize - 1) / blk_zsize))
            block_3d = (blk_xsize, blk_ysize, blk_zsize)
            grid_3d = (grid_xsize, grid_ysize, grid_zsize)
            return block_3d, grid_3d
        if mode == '2D':
            shape_x = shapes[0]
            shape_y = shapes[1]
            block_2d = (32, 32, 1)
            grid_2d = (int(round((shape_x + block_2d[0] - 1) / 32)), int(round((shape_y + block_2d[1] - 1) / 32)), 1)
            return block_2d, grid_2d
        if mode == '1D':
            shape_x = shapes[0]
            shape_y = shapes[1]
            shape_z = shapes[2]
            block_1d = (int(min(shape_x * shape_y, 1024)), 1, 1)
            grid_1d = (int((shape_x * shape_y) / block_1d[0]), int(shape_z), 1)
            return block_1d, grid_1d

    def multislice(self, bandwidth=1/3, unified_mem=False, batch_size=None, transmit=True):
        """

        :param bandwidth:
        :return:
        """
        ## TODO: when batch_size != self.num_probes, probe positions get scrambled!!!
        if batch_size == None:
            batch_size = self.num_probes
        # checks
        if isinstance(self.potential_slices, np.ndarray) is False:
            warn('Potential slices must be calculated first before calling multi_slice\n. '
                 'Call method build_potential_slices().')
            return
        if isinstance(self.psi_k, np.ndarray) is False:
            warn('Probe in k-space must be initiated first before calling multi_slice\n. '
                 'Call method build_probe().')
            return
        if isinstance(self.probe_positions, np.ndarray) is False:
            warn('Probe positions must be calculated first before calling multislice\n. '
                 'Call method generate_probe_positions().')
            return

        # load cuda kernels and prepare arguments
        shape_x = np.int32(self.sampling[1])
        shape_y = np.int32(self.sampling[0])
        num_probes = self.max_probes
        k_semi = np.float32(self.semi_ang / self.Lambda)

        # define block/grid threads
        block_3d, grid_3d = self._get_blockgrid([self.sampling[1], self.sampling[0], num_probes], mode='3D')
        block_2d, grid_2d = self._get_blockgrid([self.sampling[1], self.sampling[0], num_probes], mode='2D')
        self.print_debug('block, grid:', block_3d, grid_3d)

        # allocate memory
        # self.probes = np.empty((self.num_probes, shape_y, shape_x), dtype=np.complex64)
        self.propag = cuda.aligned_empty((int(self.sampling[0]), int(self.sampling[1])), np.complex64)
        self.vars.append(self.propag.base)
        self.propag = cuda.register_host_memory(self.propag)
        propag_d = cuda.to_device(self.propag)
        self.mask = cuda.aligned_empty((int(self.sampling[0]), int(self.sampling[1])), np.float32)
        self.vars.append(self.mask.base)
        self.mask = cuda.register_host_memory(self.mask)
        mask_d = cuda.to_device(self.mask)
        grid_steps_d = cuda.to_device(self.grid_steps.astype(np.int32))
        grid_range_d = cuda.to_device(self.grid_range.astype(np.float32))
        psi_k_d = cuda.to_device(self.psi_k)
        if unified_mem:
            self.probes = cuda.managed_empty(shape=(int(self.num_probes), int(self.sampling[0]), int(self.sampling[1])),
                                        dtype=np.complex64, mem_flags=cuda.mem_attach_flags.GLOBAL)
        else:
            # pinned memory is default
            self.probes = cuda.aligned_empty((int(self.num_probes), int(self.sampling[0]), int(self.sampling[1])), np.complex64)
            self.vars.append(self.probes.base)
            self.probes = cuda.register_host_memory(self.probes)
        ones = cuda.aligned_zeros((int(self.sampling[0]), int(self.sampling[1])), np.complex64) + 1
        self.vars.append(ones.base)
        ones = cuda.register_host_memory(ones)
        ones_d = cuda.mem_alloc(ones.nbytes)
        cuda.memcpy_htod_async(ones_d, ones, cuda.Stream())

        # grab needed kernels
        propag_func = self.kernels['propagator']
        mask_func = self.kernels['hard_aperture']
        multwise_func = self.kernels['mult_wise_c2d_re2d']

        # 1. Build propagator and bandwidth limiting mask
        mask_func(mask_d, np.float32(self.kmax), np.float32(bandwidth * self.kmax), shape_x, shape_y, block=block_2d,
                  grid=grid_2d, shared=0, stream=cuda.Stream())
        #ctx.synchronize()
        propag_func(propag_d, np.float32(self.kmax), np.float32(self.slice_t), np.float32(self.Lambda), shape_x, shape_y,
                    block=block_2d, grid=grid_2d, shared=0, stream=cuda.Stream())
        #ctx.synchronize()
        multwise_func(propag_d, mask_d, shape_x, shape_y, block=block_2d, grid=grid_2d, shared=0)
        cuda.memcpy_dtoh_async(self.propag, propag_d, cuda.Stream())
        cuda.memcpy_dtoh_async(self.mask, mask_d, cuda.Stream())

        # 1. Generate batch fft plans, create host/device pointers, and launch cuda streams
        num_phases = self.num_probes//self.max_probes
        phases = []
        self.results = []
        t = time()
        num_stream=0
        for i in range(num_phases + 1):
            phase = slice(i * self.max_probes, (i+1) * self.max_probes)
            if i == num_phases :
                phase = slice(i * self.max_probes, self.num_probes)
            phases.append(phase)
        self.print_debug('Simulation split into %d serial phases.' % len(phases))
        self.print_debug('Phases: %s ' %format(phases))
        for (i, phase) in enumerate(phases):
            t = time()
            if self.probes[phase].shape[0] == 0: break
            self.print_verbose('Starting simulation phase #%d' % i)
            num_batches = (phase.stop - phase.start + 1)//batch_size
            batches = []
            streams = []
            probes_d = []
            norm_consts = []
            plans = []
        # 2. Generate batch fft plans and create pinned memory pointers and cuda streams
            for batch_num in range(num_batches+1):
                if batch_num == num_batches:
                    slice_obj = slice(batch_num * batch_size, None)
                else:
                    slice_obj = slice(batch_num * batch_size, (batch_num + 1) * batch_size)
                batches.append(slice_obj)
                stream = cuda.Stream()
                num_stream += 1
                self.print_debug(' Launch Stream # %d' % num_stream)
                streams.append(stream)
                num_probes = np.int32(self.probe_positions[phase][slice_obj].shape[0])
                if num_probes == 0: break
                probes_d.append(cuda.mem_alloc(int(num_probes*np.prod(self.sampling)*8)))
                plans.append(skfft.Plan(self.sampling, np.complex64, np.complex64, batch=num_probes, stream=stream))
                norm_consts.append(cuda.mem_alloc(np.empty(num_probes,dtype=np.float32).nbytes))
        # 3. Propagate Beams
            for batch, stream, probe_d, plan,  norm_const in zip(batches, streams, probes_d, plans, norm_consts):
                num_probes = np.int32(self.probe_positions[phase][batch].shape[0])
                self.print_debug('batch: %s' % format(batch))
                self.__propagate_beams(num_probes, batch, probe_d, propag_d, psi_k_d, norm_const, grid_steps_d, grid_range_d,
                                     self.probes[phase][batch], plan, ones_d, stream, transmit=transmit)
            ctx.synchronize()
        # 4. clean-up
            for plan, probe_d, norm_const in zip(plans, probes_d, norm_consts):
               cufft.cufftDestroy(plan.handle)
               probe_d.free()
               norm_const.free()
               # ctx.synchronize()
               del probe_d, norm_const, plan
            ctx.synchronize()
            self.print_verbose('finished simulation phase #%d' % i)
            self.probes /= self.normalization
            # self.probes[phase][batch] = self.probes[phase][batch]/self.normalization
            sim_t = time()-t
            self.print_verbose('Propagated %d probes in %2.4f s' % (self.probe_positions[phase].shape[0], sim_t))

        if transmit:
            with catch_warnings():
                simplefilter('ignore')
                catch_warn()
                self.probes = self.probes.astype(np.float32) # discard imaginary
        # clean up device variables
#         self.vars.append(self.pot_dev_ptr) # don't de-allocate here to allow for simulations with different probe/exp. params
#         self.pot_dev_ptr.free()
        mask_d.free()
        propag_d.free()
        psi_k_d.free()
        grid_range_d.free()
        grid_steps_d.free()
        ones_d.free()
        return None

    def build_potential_slices(self, slice_thickness):
        self.slice_t = slice_thickness
        self.num_slices = np.int32(np.floor(self.dims[-1] / slice_thickness))
        tasks = [((self, slice_num), {'method': 'build_slices'}) for slice_num in range(self.num_slices)]
        processes = min(mp.cpu_count(), self.num_slices)
        chunk = np.int(np.floor(self.num_slices / processes))
        pool = mp.Pool(processes=processes)
        jobs = pool.imap(unwrap, tasks, chunksize=chunk)
        self.potential_slices = np.array([j for j in jobs])
        pool.close()
        self.print_verbose('Built %d potential slices with shape:%s pixels' % (self.potential_slices.shape[0],
                                                                  format(self.potential_slices.shape[1:])))
#         self.pot_dev_ptr = [cuda.to_device(pot_slice) for pot_slice in self.potential_slices]
        self.pot_slices_d = [cuda.to_device(pot_slice) for pot_slice in self.potential_slices]
        self.vars = []
#         print(self.potential_slices)

    def __propagate_beams(self, num_probes, batch, psi_pos_d, propag_d, psi_k_d,
                        norm_const_d, grid_steps_d, grid_range_d,
                        psi_x_pos_pin, fft_plan_probe, ones_d, stream, transmit=True):
        """
        :param batch:
        :param psi_pos_d:
        :param propag_d:
        :param psi_k_d:
        :param norm_const_d:
        :param grid_steps_d:
        :param grid_range_d:
        :param psi_x_pos_pin:
        :param fft_plan_probe:
        :param stream:
        :return:
        """
        # define block/grid threads
        block_3d, grid_3d = self._get_blockgrid([self.sampling[1], self.sampling[0], num_probes], mode='3D')
        block_1d, grid_1d = self._get_blockgrid([self.sampling[1], self.sampling[0], num_probes], mode='1D')
        self.print_debug('block, grid:', block_3d, grid_3d)

        # grab needed kernels
        probe_stack_func = self.kernels['probes_stack']
        multwise2d_stack_func = self.kernels['mult_wise_c3d_c2d']
        multwise_stack_func = self.kernels['mult_wise_c3d_c2d']
        multwise_stack_func.prepare(("P", "P", np.int32, np.float32))
#         multwise_stack_func = self.kernels['mult_wise_c3d_c3d_ind']
#         multwise_stack_func.prepare(("P", "P", np.int32, np.int32, np.float32))
        multwise2d_stack_func.prepare(("P", "P", np.int32, np.float32))
        fftshift_func = self.kernels['fftshift_2d_stack']
        # norm_const_func = self.kernels['norm_const']
        normalize_func = self.kernels['normalize']
        mod_square_func = self.kernels['mod_square_stack']

        #1. build probes
        probe_stack_func(psi_pos_d, psi_k_d, num_probes, np.float32(self.kmax), grid_steps_d, grid_range_d,
                         block=block_3d, grid=grid_3d, shared=0, stream=stream)
        #ctx.synchronize()
        cufft.cufftExecC2C(fft_plan_probe.handle, int(psi_pos_d), int(psi_pos_d), cufft.CUFFT_INVERSE)
        #ctx.synchronize()
        fftshift_func(psi_pos_d, num_probes, block=block_3d, grid=grid_3d, shared=0, stream=stream)
        #ctx.synchronize()
        ### TODO: Check why cuda normalization is off

        # 2. Propagate probes through atomic potential
        # ctx.synchronize()
        if transmit:
            for i in range(self.num_slices):
                # self.print_debug('Atomic potential slice #%d' % i)
                multwise_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, self.pot_slices_d[i], num_probes,
                                                    np.float32(1/np.prod(self.sampling)))

#                 multwise_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, self.pot_dev_ptr[i], num_probes, np.int32(i),
#                                                         np.float32(1/np.prod(self.sampling)))
#                 ctx.synchronize()
                cufft.cufftExecC2C(fft_plan_probe.handle, int(psi_pos_d), int(psi_pos_d), cufft.CUFFT_FORWARD)
#                 ctx.synchronize()
                multwise2d_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, propag_d, num_probes,
                                                        np.float32(1.0))
#                 ctx.synchronize()
                cufft.cufftExecC2C(fft_plan_probe.handle, int(psi_pos_d), int(psi_pos_d), cufft.CUFFT_INVERSE)
#                 ctx.synchronize()

            multwise2d_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, ones_d, num_probes,
                                                    np.float32(1 / (np.prod(self.sampling))))
            cufft.cufftExecC2C(fft_plan_probe.handle, int(psi_pos_d), int(psi_pos_d), cufft.CUFFT_FORWARD)
            multwise2d_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, ones_d, num_probes,
                                                    np.float32(np.sqrt(np.prod(self.sampling))))
            mod_square_func(psi_pos_d, num_probes, block=block_3d, grid=grid_3d, stream=stream)
        # ctx.synchronize()
        cuda.memcpy_dtoh_async(psi_x_pos_pin, psi_pos_d, stream=stream)

class MSAMPI(MSAGPU):
    def __init__(self, *args, **kwargs):
        global comm
        comm = MPI.COMM_WORLD
        self.size = comm.Get_size()
        self.rank = comm.Get_rank()
        super(MSAMPI, self).__init__(*args, **kwargs)

    def print_rank(self, *args, **kwargs):
        if self.rank == 0:
            print(*args, **kwargs)

    def print_verbose(self, *args, **kwargs):
        if self.verbose:
            self.print_rank(*args, **kwargs)

    def generate_probe_positions(self, *args, **kwargs):
        super(MSAMPI, self).generate_probe_positions(*args, **kwargs)
        data_parts = []
        part = self.probe_positions.shape[0] // self.size
        self.total_num_probes = self.probe_positions.shape[0]
        for rank in range(self.size):
            slice_obj = slice(rank * part, (rank + 1)* part)
            if rank == self.size - 1:
                slice_obj = slice( rank * part, None)
            data_parts.append(slice_obj)
        self.data_parts = data_parts
        if self.rank == 0:
            data = [self.probe_positions[part] for part in data_parts]
            self.data_size = np.sum([dat.shape[0] for dat in data])
        else:
            data = None
        data = comm.scatter(data, root=0)
        self.probe_positions = data
        print("rank %d: my # of probe_positions %d" % (self.rank,
        self.probe_positions.shape[0]))
        self.num_probes = np.int32(self.probe_positions.shape[0])

    def multislice(self, *args, **kwargs):
        try:
            # h5_write= kwargs.pop('h5_write')
            h5_file = kwargs.pop('h5_write')
            h5_write = isinstance(h5_file, h5py.File)
        except KeyError:
            h5_write = False
        super(MSAMPI, self).multislice(*args, **kwargs)
        if h5_write:
            self.print_rank('writing to h5 file...')
            dset = h5_file.create_dataset('4D_CBED', (self.total_num_probes,
                                self.sampling[0], self.sampling[1]), dtype=self.probes.dtype)
            # Workaround the large parallel I/O limitation
            ### TODO: (low priority) test if doing below in chunks over indices < 2GB speeds up the writing.
            for i, p in enumerate(range(*self.data_parts[self.rank].indices(self.total_num_probes))):
                dset[p, :, :] = self.probes[i, :, :]
            self.print_rank('finished writing h5 file.')
        else:
            receive_buff = None
            self.print_verbose('rank %d: shape of calculated probes = %s' %(self.rank, format(self.probes.shape)))
            if self.rank == 0:
                buff_shape = (self.total_num_probes, self.sampling[0], self.sampling[1])
                self.print_rank('receive buffer shape: %s' %format(buff_shape))
                receive_buff = np.empty(buff_shape, dtype=self.probes.dtype)
            # Calculate buffer size for a vector gather op
            count = []
            for data_part in self.data_parts:
                start = data_part.start
                if start is None:
                    start = 0
                stop = data_part.stop
                if stop is None:
                    stop = self.total_num_probes
                count.append(stop - start)
            count = np.array(count)
            displ = np.cumsum(count)
            displ = np.append(displ[::-1],0)[::-1]
            xy_offset = np.prod(self.sampling)
            count *= xy_offset
            count = count.astype(np.int64)
            displ *= xy_offset
            displ = displ.astype(np.int64)
            self.print_verbose("rank %d: my # probes: %s %s, count: %d, displ: %d" %(self.rank, format(self.probes.shape), self.probes.dtype, count[self.rank], displ[self.rank]))
            # gather
            # comm.Gatherv(self.probes, [receive_buff, count, displ[:-1], MPI.C_FLOAT_COMPLEX], root=0)
            comm.Gatherv(self.probes, [receive_buff, count, displ[:-1], MPI.FLOAT], root=0)

            if self.rank == 0: self.print_rank('Gathered results of shape: %s' %format(receive_buff.shape))

max_stars_repo_name: nlaanait/namsa

max_issues_repo_name: nlaanait/namsa

max_forks_repo_name: nlaanait/namsa

reason: import numpy,from scipy,import pycuda,from pycuda


text:
# Licensed under an MIT open source license - see LICENSE

from ..length import skeleton_length
from ..pixel_ident import is_blockpoint, is_tpoint

import numpy as np

def test_length1():

    test_skel1 = np.zeros((10, 10))

    test_skel1[1, 1:3] = 1
    test_skel1[2, 3] = 1
    test_skel1[3, 4:6] = 1

    length1 = skeleton_length(test_skel1)

    assert length1 == 2 + np.sqrt(2) * 2

def test_length2():

    test_skel2 = np.eye(10)

    length2 = skeleton_length(test_skel2)

    assert length2 == 9 * np.sqrt(2)

def test_length3():

    test_skel3 = np.zeros((10, 10))

    test_skel3[:, 5] = 1

    length3 = skeleton_length(test_skel3)

    assert length3 == 9

def test_length4():

    test_skel4 = np.zeros((12, 12))

    posns = \
        [(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 11, 11, 10, 9, 9),
         (3, 2, 1, 0, 1, 2, 3, 3, 3, 4, 5, 6, 7, 8, 9, 10, 11)]

    test_skel4[posns] = 1

    length4 = skeleton_length(test_skel4)

    assert length4 == 11 * np.sqrt(2) + 5

def test_length5():

    test_skel5 = np.zeros((9, 12))

    posns = \
        [(8, 8, 7, 7, 6, 5, 4, 3, 2, 1, 0, 0, 1, 1, 2, 3, 4, 5),
         (0, 1, 2, 3, 4, 4, 4, 5, 5, 5, 6, 7, 8, 9, 10, 10, 11, 11)]

    test_skel5[posns] = 1

    length5 = skeleton_length(test_skel5)

    assert length5 == 10 + 7 * np.sqrt(2)

def test_length6():

    test_skel6 = np.zeros((4, 4))

    test_skel6[1, 0:2] = 1
    test_skel6[2:4, 2] = 1

    length6 = skeleton_length(test_skel6)

    assert length6 == 2 + np.sqrt(2)

def test_isblockpt():

    # test four different configurations, and once with an added connection

    # 0 1 6
    vallist = [1, 1, 0, 0, 0, 0, 0, 1]

    assert is_blockpoint(vallist)

    vallist = [1, 1, 0, 0, 1, 0, 0, 1]

    assert is_blockpoint(vallist)

    # 1 2 3
    vallist = [0, 1, 1, 1, 0, 0, 0, 0]

    assert is_blockpoint(vallist)

    vallist = [0, 1, 1, 1, 0, 1, 0, 0]

    assert is_blockpoint(vallist)

    # 3 4 5
    vallist = [0, 0, 0, 1, 1, 1, 0, 0, 0]

    assert is_blockpoint(vallist)

    vallist = [1, 0, 0, 1, 1, 1, 0, 0, 0]

    assert is_blockpoint(vallist)

    # 5 6 7
    vallist = [0, 0, 0, 0, 0, 1, 1, 1]

    assert is_blockpoint(vallist)

    vallist = [0, 1, 0, 0, 0, 1, 1, 1]

    assert is_blockpoint(vallist)

def test_isTpoint():

    # 0 6 7
    vallist = [1, 0, 0, 0, 0, 0, 1, 1]

    assert is_tpoint(vallist)

    vallist = [1, 0, 0, 1, 0, 0, 1, 1]

    assert is_tpoint(vallist)

    # 0 1 2
    vallist = [1, 1, 1, 0, 0, 0, 0, 0]

    assert is_tpoint(vallist)

    vallist = [1, 1, 1, 0, 1, 0, 0, 0]

    assert is_tpoint(vallist)

    # 2 3 4
    vallist = [0, 0, 1, 1, 1, 0, 0, 0, 0]

    assert is_tpoint(vallist)

    vallist = [0, 0, 1, 1, 1, 0, 0, 1, 0]

    assert is_tpoint(vallist)

    # 4 5 6
    vallist = [0, 0, 0, 0, 1, 1, 1, 0]

    assert is_tpoint(vallist)

    vallist = [1, 0, 0, 0, 1, 1, 1, 0]

    assert is_tpoint(vallist)

ext: py

reason: import numpy


text:
from dask.distributed import Client, LocalCluster
import numpy as np

def dask_eval(xx, eval_func, client, serial=False):
    ''' Performs parallel evaluation across dask workers'''
    objs, descs, phenos = [], [], []
    if serial:
        for x in xx:
            obj, desc, pheno = eval_func(x)
            objs.append(obj)
            descs.append(desc)
            phenos.append(pheno)
    else:
        futures = client.map(lambda x: eval_func(x), xx)
        results = client.gather(futures)
        for obj, desc, pheno in results:
            objs.append(obj)
            descs.append(desc)
            phenos.append(pheno)

    objs = np.hstack(objs)
    descs = np.vstack(descs)

    return objs, descs, phenos

def create_dask_client(n_workers):
    ''' Creats local cluster of dask workers'''
    cluster = LocalCluster(
    processes=True,  # Each worker is a process.
    n_workers=n_workers,  # Create this many worker processes.
    threads_per_worker=1,  # Each worker process is single-threaded.
    )
    return Client(cluster)

max_stars_repo_name: agaier/bbq

max_issues_repo_name: agaier/bbq

max_forks_repo_name: agaier/bbq


text:
from __future__ import print_function
import itertools
import sys, os
import inspect
import numpy as np
from joblib import load as joblib_load
from .resultsManager import ResultsManager
from .classifiers.randomForest import  predictMethod

from Config import Configuration

class ComplexPredictor(Configuration):
  '''
    This class is used to predict new pdbs once models have already been computed.
  '''
  def __init__(self, stepName, savedModelsPath=None):
    '''

      @param stepName: str. Must startswith seq_train or struct or mixed (seq_train, mixed_2, structX, seq_train1... are also valid)
      @param savedModelsPath: str. A path to the directory where models have been saved. If None,
                                   it will used the path indicated in Config
    '''
    Configuration.__init__(self)

    self.stepName= stepName
    if not savedModelsPath is None:
      self.savedModelsPath= savedModelsPath

    self.model=None
    print(stepName)
    for fname in os.listdir(self.savedModelsPath):
      if fname.endswith(stepName):
        print("Loading model %s"%(fname))
        self.model= joblib_load(os.path.join(self.savedModelsPath, fname))
    assert not self.model is None, "Error, there is no valid model in %s for step %s"%(self.savedModelsPath, self.stepName)

  def predictOneComplex(self, complexCodifiedObject, outName, isLastStep=False):
    '''
      computes pairwise RR interaction and binding-site predictions for the codified complex passed as argument

      @param complexCodifiedObject: A codifyComplexes.ComplexCodified object that represents a protein protein
                                    interaction
      @param outName: str.  The name that RR interaction results file will have. Ligand binding-site results
                            name will be outName+'lig' and Receptor binding-site results
                            name will be outName+'rec'
      @param outName: boolean. True if no more feedback steps will be done
    '''
    data_d, data_t= complexCodifiedObject.getData()
    data_ids= complexCodifiedObject.getIds()
    prefix= complexCodifiedObject.getPrefix()
    prob_predictionsDir= predictMethod(self.model, data_d)
    prob_predictionsTrans= predictMethod(self.model, data_t)
    resultObj= ResultsManager(prefix, prob_predictionsDir, prob_predictionsTrans, data_ids)
    p,l,r= resultObj.getResultsPdDf("p"), resultObj.getResultsPdDf("l"), resultObj.getResultsPdDf("r")
    resultObj.writeResults(outName, isLastStep)
    return  outName, (p,l,r)

max_stars_repo_name: bioinsilico/BIPSPI

max_issues_repo_name: bioinsilico/BIPSPI

max_forks_repo_name: bioinsilico/BIPSPI


text:
""" ConnectedComponents.py """

import sys
import time
import numpy as np
import networkx as nx
from ..exrpc.server import FrovedisServer
from ..exrpc import rpclib
from .graph import Graph
from .g_validate import validate_graph

def connected_components(G, opt_level=2, hyb_threshold=0.4,
                         verbose=0,
                         print_summary=False,
                         print_limit=5):
    """
    DESC:   Computes connected components of a graph
    PARAM:  G: frovedis/networkx graph object
            opt_level: int (0, 1 or 2) (default: 2)
            hyb_threshold: double (0.0 ~ 1.0) (default: 0.4)
            verbose: int (0 or 1) (default: 0)
            print_summary: whether to print summary of connected components
            print_limit: the maximum number of nodes info to be printed
    RETURN: A dictionary with keys as root-nodeid for each component,
            and values as list of pairs of nodeid with its distance
            from root of the component to which the node belongs
    """
    G, inp_movable = validate_graph(G)
    nvert = G.num_vertices
    (host, port) = FrovedisServer.getServerInstance()
    # graph.py: node data is loaded as int64
    nodes_dist = np.empty(nvert, dtype=np.int64) # actually levels (I-type)
    nodes_in_which_cc = np.empty(nvert, dtype=np.int64)

    stime = time.time()
    root_with_cc_count = rpclib.call_frovedis_cc(host, port,\
                                    G.get(), nodes_in_which_cc,\
                                    nodes_dist, nvert,\
                                    opt_level, hyb_threshold)
    excpt = rpclib.check_server_exception()
    if excpt["status"]:
        raise RuntimeError(excpt["info"])
    if inp_movable:
        G.release()
    etime = time.time()
    if verbose:
        print("cc computation time: %.3f sec." % (etime - stime))

    stime = time.time()
    if print_summary:
        show_cc_summary(root_with_cc_count, nodes_in_which_cc, \
                        nodes_dist, nvert, print_limit)
    ret = {root_with_cc_count[i] : [] \
           for i in range(0, len(root_with_cc_count), 2)}
    for i in range(0, nvert):
        cc_root = nodes_in_which_cc[i]
        if cc_root != sys.maxsize:
            ret[cc_root].append(i + 1)
    etime = time.time()
    if verbose:
        print("cc res conversion time: %.3f sec." % (etime - stime))
    for i in sorted(ret.values()):
        yield set(i)

def show_cc_summary(root_with_cc_count,
                    nodes_in_which_cc, nodes_dist,
                    num_vertices, print_limit):
    """Connected Component Summary"""
    num_cc = len(root_with_cc_count) // 2
    print("Number of connected components: %d" % num_cc)

    num_cc_printed = min(print_limit, num_cc)
    print("Root with its count of nodes in each connected component: " \
          + "(root_id:count)")
    for i in range(0, 2 * num_cc_printed, 2):
        print("%d:%d  " % (root_with_cc_count[i], \
                           root_with_cc_count[i + 1]))
    if num_cc > print_limit:
        print("...")

    print("Nodes in which cc: (node_id:root_id)")
    count = 1
    for i in range(num_vertices):
        if nodes_in_which_cc[i] != sys.maxsize:
            print("%d:%d  " % (i+1, nodes_in_which_cc[i]))
            count = count + 1
        if count > print_limit:
            print("...")
            break

    print("Nodes dist: (node:level_from_root)")
    count = 1
    for i in range(num_vertices):
        if nodes_dist[i] != sys.maxsize:
            print("%d:%d  " % (i+1, nodes_dist[i]))
            count = count + 1
        if count > print_limit:
            print("...")
            break

max_stars_repo_name: XpressAI/frovedis

max_issues_repo_name: XpressAI/frovedis

max_forks_repo_name: XpressAI/frovedis

reason: import numpy,import networkx


text:
import sys
import os
import time

BASE_PATH = os.path.abspath(os.path.join(os.getcwd()))
sys.path.append(BASE_PATH)
import pickle
import torch
import torch.optim as optim
import data_generator
import tools
from args import read_args
import numpy as np
import random
torch.set_num_threads(2)
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

class model_class(object):
	def __init__(self, args):
		super(model_class, self).__init__()
		self.args = args
		self.gpu = args.cuda

		input_data = data_generator.input_data(args=self.args)
		self.input_data = input_data

		if self.args.train_test_label == 2:
			input_data.het_walk_restart()
			print("neighbor set generation finish")

		feature_list = [input_data.p_title_embed,\
		input_data.p_v_net_embed, input_data.p_a_net_embed, input_data.p_ref_net_embed,\
		input_data.p_net_embed, input_data.a_net_embed, input_data.a_text_embed,\
		input_data.v_net_embed, input_data.v_text_embed]

		for i in range(len(feature_list)):
			feature_list[i] = torch.from_numpy(np.array(feature_list[i])).float()

		if self.gpu:
			for i in range(len(feature_list)):
				feature_list[i] = feature_list[i].cuda()

		a_neigh_list_train = input_data.a_neigh_list_train
		p_neigh_list_train = input_data.p_neigh_list_train
		v_neigh_list_train = input_data.v_neigh_list_train

		a_train_id_list = input_data.a_train_id_list
		p_train_id_list = input_data.p_train_id_list
		v_train_id_list = input_data.v_train_id_list

		self.model = tools.HetAgg(args, feature_list, a_neigh_list_train, p_neigh_list_train, v_neigh_list_train, \
								  a_train_id_list, p_train_id_list, v_train_id_list)

		if self.gpu:
			self.model.cuda()
		self.parameters = filter(lambda p: p.requires_grad, self.model.parameters())
		self.optim = optim.Adam(self.parameters, lr=self.args.lr, weight_decay = args.weight_decay)
		self.model.init_weights()

	def model_train(self):
		print('Model Training')
		self.model.train()
		mini_batch_s = self.args.mini_batch_s
		embed_d = self.args.embed_d

		for iter_i in range(self.args.train_iter_n):
			# print('iteration ' + str(iter_i) + ' ...')
			start_time = time.time()
			triple_list = self.input_data.sample_het_walk_triple()

			min_len = 1e10
			for ii in range(len(triple_list)):
				if len(triple_list[ii]) < min_len:
					min_len = len(triple_list[ii])
			batch_n = int(min_len / mini_batch_s)

			loss_list = []
			for k in range(batch_n):
				c_out = torch.zeros([len(triple_list), mini_batch_s, embed_d])
				p_out = torch.zeros([len(triple_list), mini_batch_s, embed_d])
				n_out = torch.zeros([len(triple_list), mini_batch_s, embed_d])

				for triple_index in range(len(triple_list)):
					triple_list_temp = triple_list[triple_index]
					triple_list_batch = triple_list_temp[k * mini_batch_s : (k + 1) * mini_batch_s] #batch

					c_out_temp, p_out_temp, n_out_temp = self.model(triple_list_batch, triple_index)

					c_out[triple_index] = c_out_temp
					p_out[triple_index] = p_out_temp
					n_out[triple_index] = n_out_temp

				loss = tools.cross_entropy_loss(c_out, p_out, n_out, embed_d)

				self.optim.zero_grad()
				loss.backward()
				self.optim.step()

				# if (k+1) % 600 == 0:
				# 	print("loss: " + str(loss))
				loss_list.append(loss)

			train_loss = sum(loss_list) / len(loss_list)

			if (iter_i+1) % args.save_model_freq == 0:
				model_save_path = 'model_save/'
				if not os.path.isdir(model_save_path):
					os.mkdir(model_save_path)
				save_path = model_save_path + "SIHDGNN_" + str(iter_i) + ".pt"
				torch.save(self.model.state_dict(), save_path)
				print('Model saved at {}'.format(save_path))
				triple_index = 9
				_1,_2,_3 = self.model([], triple_index)

			print('Iteration: {} \t Train_Loss: {:.3f} \t Time: {:.3f}s' \
				  .format(iter_i, train_loss, time.time() - start_time))

if __name__ == '__main__':
	args = read_args()
	# print("------arguments-------")
	# for k, v in vars(args).items():
	# 	print(k + ': ' + str(v))

	random.seed(args.random_seed)
	np.random.seed(args.random_seed)
	torch.manual_seed(args.random_seed)
	torch.cuda.manual_seed_all(args.random_seed)

	model_object = model_class(args)

	if args.train_test_label == 0:
		model_object.model_train()


text:
import time
import logging
import warnings

import numpy as np
import xarray as xr
from dask.diagnostics import ProgressBar

from wind_repower_usa.config import MONTHS
from wind_repower_usa.load_data import load_turbines, load_wind_velocity
from wind_repower_usa.load_data import load_wind_speed
from wind_repower_usa.turbine_models import ge15_77

def calc_wind_speed_at_turbines(wind_velocity, turbines):
    # interpolate at turbine locations
    wind_velocity_at_turbines = wind_velocity.interp(
        longitude=xr.DataArray(turbines.xlong.values, dims='turbines'),
        latitude=xr.DataArray(turbines.ylat.values, dims='turbines'),
        method='linear')

    # velocity --> speed
    wind_speed = (wind_velocity_at_turbines.u100**2
                  + wind_velocity_at_turbines.v100**2)**0.5

    return wind_speed

def calc_simulated_energy(wind_speed, turbines, power_curve=None, sum_along='turbines',
                          capacity_scaling=True, only_built_turbines=True):
    """Estimate generated energy using wind data and turbine data.

    Parameters
    ----------
    wind_speed : xr.DataArray
        see calc_wind_speed_at_turbines()
    turbines : xr.DataSet
        see load_turbines()
    power_curve : callable
        a function mapping wind speed to power
    sum_along : str
        sum along turbines or time or emtpy string
    capacity_scaling : bool
        scale power curve to capacity for each turbine (if available)
    only_built_turbines : bool
        calculate energy only for time stamps where commission year is older

    Returns
    -------
    simulated_energy_gwh : xr.DataArray
        Simulated energy per month [GWh], dims = (time, turbines)

    FIXME this modifies the input wind_speed variable! Very dangerous but unclear if solvable
     without too much memory consumption (via copying)

    """
    if power_curve is None:
        power_curve = ge15_77.power_curve

    # this outputs a deprecation warning, see https://github.com/pydata/xarray/issues/2928
    # TODO probably not the best idea to have this here, since it modifies gloabl behavior at
    #  runtime, but where else to put it? pytest ignores warnings.catch_warnings()...
    warnings.filterwarnings('ignore', 'The da.atop function has moved to da.blockwise')

    # TODO this is a bit scary, when does parallelized not work? Which dtype?
    simulated_energy = xr.apply_ufunc(power_curve, wind_speed,
                                      dask='parallelized',
                                      output_dtypes=[np.float64])

    simulated_energy = simulated_energy.assign_coords(turbines=turbines.turbines)

    if only_built_turbines:
        # TODO all turbines where year = NaN will be removed that way... :-/
        # this is the beginning of the year the turbine has been commissioned
        building_dates = turbines.p_year.astype(int).astype(str).astype(np.datetime64)

        nanosecs_of_year = (simulated_energy.time - building_dates).astype(np.float)
        proportion_of_year = nanosecs_of_year / (365.25 * 24 * 60 * 60 * 1e9)

        # comparing objects with dim "time" and dim "turbines" results in (time, turbines)
        building_this_year = simulated_energy.time.dt.year == turbines.p_year

        simulated_energy = simulated_energy.where(~building_this_year,
                                                  simulated_energy * proportion_of_year)

        already_built = simulated_energy.time.dt.year >= turbines.p_year
        simulated_energy = simulated_energy.where(already_built, 0)

        # Uargh... there is a weired memory leak somewhere, this seems to help a bit at least... :-/
        del nanosecs_of_year
        del proportion_of_year
        del building_this_year
        del building_dates
        del already_built

    if capacity_scaling:
        # FIXME this should use turbine_model.capacity_mw not 1500!
        simulated_energy *= (turbines.t_cap / 1500.).fillna(1.)

    # inspired by:
    # http://xarray.pydata.org/en/stable/examples/weather-data.html#monthly-averaging

    simulated_energy = simulated_energy.sortby('time') * 1e-6

    if sum_along:
        simulated_energy = simulated_energy.sum(dim=sum_along)
    if sum_along == 'turbines':
        simulated_energy = simulated_energy.resample(time='1MS').sum()

    # Does not work for multiple years:
    # simulated_energy = simulated_energy.sum(dim='turbines').groupby('time.month').sum() * 1e-6

    with ProgressBar():
        simulated_energy_gwh = simulated_energy.compute()

    if sum_along == 'turbines':
        simulated_energy_gwh.name = "Simulated energy per month [GWh]"
    elif sum_along == 'time':
        simulated_energy_gwh.name = "Simulated energy"  # TODO unit depends on time range?

    return simulated_energy_gwh

def calc_simulated_energy_years(years, turbines=None, power_curve=None, capacity_scaling=True,
                                only_built_turbines=True):
    """Load wind speed data from processed files and calculate energy in a loop to avoid large
    data in memory for the parts where dask chunks are not correctly working yet.

    Parameters
    ----------
    years : iterable
        eg. range(2004, 2017)
    turbines : xr.DataSet
        as returned by load_turbines()
    power_curve : callable
        see calc_simulated_energy()

    Returns
    -------
    xr.DataArray

    """
    eta = False
    simulated_energy_gwh = []
    if turbines is None:
        turbines = load_turbines()

    with ProgressBar():
        for year in years:
            for month in range(1, 13):
                t0 = time.time()
                logging.info("Calculating {}-{:02d}...".format(year, month))
                wind_speed = load_wind_speed(year, month)
                simulated_energy_gwh.append(
                    calc_simulated_energy(
                        wind_speed=wind_speed,
                        turbines=turbines,
                        power_curve=power_curve,
                        capacity_scaling=capacity_scaling,
                        only_built_turbines=only_built_turbines
                    )
                )

                if not eta:
                    logging.info("ETA: %s seconds", (time.time() - t0) * 12 * len(years))
                    eta = True

    simulated_energy_gwh = xr.concat(simulated_energy_gwh, dim='time')
    return simulated_energy_gwh

def calc_bounding_box_usa(turbines, extension=1.):
    # Bounding box can be also manually selected:
    #   https://boundingbox.klokantech.com/

    # assert -180 <= long <= 180, -90 <= lat <= 90
    # FIXME need +180 modulo 360!
    north = turbines.ylat.values.max() + extension
    west = turbines.xlong.values.min() - extension
    south = turbines.ylat.values.min() - extension
    east = turbines.xlong.values.max() + extension

    return north, west, south, east

def calc_wind_speed_probablity(num_samples=200, year=2016, bins=150):
    """Take random wind speed samples in given years at all wind speed
    locations and return the probability density function.
    Returns
    -------
    wind_speed_probability : xr.DataArray
        wind_speed in m/s vs probability

    """
    # TODO code duplication with data_exploration notebook
    wind_speed = xr.concat([load_wind_speed(year, month)
                            for month in MONTHS], dim='time')

    # FIXME it would be choose in 2d over turbines/time
    idcs_time = np.random.choice(wind_speed.time.shape[0],
                                 size=num_samples)
    idcs_turbines = np.random.choice(wind_speed.turbines.shape[0],
                                     size=num_samples)

    wind_speed_values = wind_speed.isel(time=idcs_time,
                                        turbines=idcs_turbines).values

    probabilities_hist, wind_speed_hist = np.histogram(
        wind_speed_values.flatten(), density=True, bins=bins)

    # there are more edges of bins than values --> picket fence error
    wind_speed_hist = wind_speed_hist[1:] - np.diff(wind_speed_hist)/2.

    wind_speed_probality = xr.DataArray(probabilities_hist, dims='wind_speed',
                                        coords={'wind_speed': wind_speed_hist})
    return wind_speed_probality

def calc_mean_wind_speed(years, sample_size=200):
    wind_velocity = load_wind_velocity(years, MONTHS)
    idcs_time = np.random.choice(wind_velocity.time.shape[0],
                                 size=sample_size, replace=False)

    wind_speed_mean = ((wind_velocity.v10**2 +
                        wind_velocity.u10**2)**0.5).isel(time=idcs_time).mean(dim='time')

    return wind_speed_mean

max_stars_repo_name: INWE-BOKU/repowering-potential-usa

max_issues_repo_name: inwe-boku/wind-repowering-usa

max_forks_repo_name: inwe-boku/wind-repowering-usa


text:
import logging
import os.path
import matplotlib.pyplot as plt
import scipy.stats as stats
import numpy as np

from pathlib import Path
from datetime import datetime
from fancylog import fancylog
import opendirection.run.run as run

def main():
    start_time = datetime.now()
    config_path = os.path.join(str(Path(__file__).parents[1]), "options")

    args, options, config = run.setup(config_path)
    logging.info("Running full analysis of each condition")
    conditions, total_df = run.analysis(args, options, config)

    if options.stability:
        conditions, _, _ = run_stability_analyses(
            args, options, config, conditions
        )

    if args.save:
        save_dfs = {}
        logging.info("Saving results")
        for condition in conditions:
            save_dfs[condition.name] = run.save(condition, args.output_dir)
        if args.summary_paths is not None:
            logging.info("saving summary csvs")
            for summary_path in args.summary_paths:
                run.summary(summary_path, save_dfs, args.output_dir)

    logging.info(
        "Finished calculations. Total time taken: %s",
        datetime.now() - start_time,
    )

    fancylog.disable_logging()

    if args.plot:
        logging.info("PLOTTING:")
        for condition in conditions:
            run.plot(
                condition.all_data,
                condition.cell_specific_data,
                condition.cell_specific_stats,
                condition.hd,
                condition.ahv,
                condition.velocity,
                condition.place,
                condition.colors,
                config.camera_frames_per_sec,
                options,
                condition_name=condition.name,
            )
    plt.show()

def run_stability_analyses(args, options, config, conditions):
    logging.info("Running analysis on first half of each condition")

    # cell list is the same for all conditions
    conditions_first, _ = run.analysis(
        args,
        options,
        config,
        stability="first",
        cell_list=conditions[0].cell_list,
    )

    logging.info("Running analysis on second half of each condition")
    conditions_last, _ = run.analysis(
        args,
        options,
        config,
        stability="last",
        cell_list=conditions[0].cell_list,
    )

    (
        conditions,
        condition_r_velocity_distributions,
        condition_r_ahv_distributions,
    ) = add_stability_indices(conditions, conditions_first, conditions_last)
    return (
        conditions,
        condition_r_velocity_distributions,
        condition_r_ahv_distributions,
    )

def stability_calcs(conditions_first, conditions_last):
    condition_velocity_correlations = []
    condition_ahv_correlations = []

    condition_r_velocity_distributions = []
    condition_r_ahv_distributions = []

    condition_velocity_null_percentiles = []
    condition_ahv_null_percentiles = []

    for idx, condition in enumerate(conditions_first):

        velocity_correlations = []
        ahv_correlations = []

        r_velocity_distributions = []
        r_ahv_distributions = []

        velocity_null_percentiles = []
        ahv_null_percentiles = []
        for cell_id, cell in enumerate(condition.cell_list):
            _r_velocity_distributions = correlate_null_distributions(
                conditions_first[idx]
                .cell_specific_stats[cell_id]
                .velocity._VelocityStats__shuffled_binned_data,
                conditions_last[idx]
                .cell_specific_stats[cell_id]
                .velocity._VelocityStats__shuffled_binned_data,
            )

            min_length = min(
                len(
                    conditions_first[
                        idx
                    ].cell_specific_data.velocity_cell_spikes_freq[cell_id]
                ),
                len(
                    conditions_last[
                        idx
                    ].cell_specific_data.velocity_cell_spikes_freq[cell_id]
                ),
            )
            r, p = stats.pearsonr(
                conditions_first[
                    idx
                ].cell_specific_data.velocity_cell_spikes_freq[cell_id][
                    :min_length
                ],
                conditions_last[
                    idx
                ].cell_specific_data.velocity_cell_spikes_freq[cell_id][
                    :min_length
                ],
            )

            velocity_correlations.append((r, p))
            velocity_null_percentile = stats.percentileofscore(
                abs(_r_velocity_distributions), abs(r)
            )
            velocity_null_percentiles.append(velocity_null_percentile)
            r_velocity_distributions.append(_r_velocity_distributions)

            _r_ahv_distributions = correlate_null_distributions(
                conditions_first[idx]
                .cell_specific_stats[cell_id]
                .ahv._AHVStats__shuffled_binned_data,
                conditions_last[idx]
                .cell_specific_stats[cell_id]
                .ahv._AHVStats__shuffled_binned_data,
            )

            r, p = stats.pearsonr(
                conditions_first[idx].cell_specific_data.ahv_cell_spikes_freq[
                    cell_id
                ],
                conditions_last[idx].cell_specific_data.ahv_cell_spikes_freq[
                    cell_id
                ],
            )

            ahv_correlations.append((r, p))
            ahv_null_percentile = stats.percentileofscore(
                abs(_r_ahv_distributions), abs(r)
            )
            ahv_null_percentiles.append(ahv_null_percentile)
            r_ahv_distributions.append(_r_ahv_distributions)

        condition_velocity_correlations.append(velocity_correlations)
        condition_ahv_correlations.append(ahv_correlations)

        condition_r_velocity_distributions.append(r_velocity_distributions)
        condition_r_ahv_distributions.append(r_ahv_distributions)

        condition_velocity_null_percentiles.append(velocity_null_percentiles)
        condition_ahv_null_percentiles.append(ahv_null_percentiles)
    return (
        condition_velocity_correlations,
        condition_ahv_correlations,
        condition_velocity_null_percentiles,
        condition_ahv_null_percentiles,
        condition_r_velocity_distributions,
        condition_r_ahv_distributions,
    )

def correlate_null_distributions(
    shuffled_tunings_first, shuffled_tunings_last
):
    r_values = []
    for first, last in zip(shuffled_tunings_first, shuffled_tunings_last):
        min_length = min(len(first), len(last))
        r_values.append(
            stats.pearsonr(first[:min_length], last[:min_length])[0]
        )
    return np.array(r_values)

def add_stability_indices(conditions, conditions_first, conditions_last):
    (
        condition_velocity_correlations,
        condition_ahv_correlations,
        condition_velocity_null_percentiles,
        condition_ahv_null_percentiles,
        condition_r_velocity_distributions,
        condition_r_ahv_distributions,
    ) = stability_calcs(conditions_first, conditions_last)

    for idx, condition in enumerate(conditions):
        # r is first output of stats.pearsonr
        for cell_id, cell in enumerate(condition.cell_specific_stats):
            # ahv
            cell.ahv.ahv_stability_index = condition_velocity_correlations[
                idx
            ][cell_id][0]

            cell.ahv.ahv_pearson_r_first_half_neg = (
                conditions_first[idx]
                .cell_specific_stats[cell_id]
                .ahv.ahv_pearson_r_neg
            )
            cell.ahv.ahv_pearson_r_first_half_pos = (
                conditions_first[idx]
                .cell_specific_stats[cell_id]
                .ahv.ahv_pearson_r_pos
            )

            cell.ahv.ahv_pearson_r_second_half_neg = (
                conditions_last[idx]
                .cell_specific_stats[cell_id]
                .ahv.ahv_pearson_r_neg
            )
            cell.ahv.ahv_pearson_r_second_half_pos = (
                conditions_last[idx]
                .cell_specific_stats[cell_id]
                .ahv.ahv_pearson_r_pos
            )

            cell.ahv.ahv_r_percentile_first_half_neg = (
                conditions_first[idx]
                .cell_specific_stats[cell_id]
                .ahv.pearson_neg_percentile
            )
            cell.ahv.ahv_r_percentile_first_half_pos = (
                conditions_first[idx]
                .cell_specific_stats[cell_id]
                .ahv.pearson_pos_percentile
            )

            cell.ahv.ahv_r_percentile_second_half_neg = (
                conditions_last[idx]
                .cell_specific_stats[cell_id]
                .ahv.pearson_neg_percentile
            )
            cell.ahv.ahv_r_percentile_second_half_pos = (
                conditions_last[idx]
                .cell_specific_stats[cell_id]
                .ahv.pearson_pos_percentile
            )

            cell.ahv.ahv_null_correlation_percentile = condition_ahv_null_percentiles[
                idx
            ][
                cell_id
            ]

            # velocity
            cell.velocity.velocity_stability_index = condition_ahv_correlations[
                idx
            ][
                cell_id
            ][
                0
            ]

            cell.velocity.velocity_r_percentile_first_half = (
                conditions_first[idx]
                .cell_specific_stats[cell_id]
                .velocity.pearson_percentile
            )
            cell.velocity.velocity_r_percentile_second_half = (
                conditions_last[idx]
                .cell_specific_stats[cell_id]
                .velocity.pearson_percentile
            )

            cell.velocity.velocity_pearson_r_first_half = (
                conditions_first[idx]
                .cell_specific_stats[cell_id]
                .velocity.velocity_pearson_r
            )

            cell.velocity.velocity_pearson_r_second_half = (
                conditions_last[idx]
                .cell_specific_stats[cell_id]
                .velocity.velocity_pearson_r
            )

            cell.velocity.velocity_null_correlation_percentile = condition_velocity_null_percentiles[
                idx
            ][
                cell_id
            ]

    return (
        conditions,
        condition_r_velocity_distributions,
        condition_r_ahv_distributions,
    )

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

max_stars_repo_name: adamltyson/opendirection

max_issues_repo_name: adamltyson/opendirection

max_forks_repo_name: adamltyson/opendirection

reason: import numpy,import scipy