text:
from __future__ import division, unicode_literals, absolute_import
import numpy as np

from .strain import Series, filtering

def evaluate_psd(strain, dt, subseglen, overlap_fraction = 0.99):
    """
        Compute the power spectral density by
        Welch's average periodogram method.
        The vector strain is divided into sebsegments
        with subseglen duration.

        Arguments:

            - strain            : noise data strain array
            - dt                : time step of the input strain
            - subseglen         : seglen of subsegment
            - overlap_fraction  : fraction of overlapping elements, default 0.99

        Return:
            - freqs : frequency axis
            - psd   : power spectral density
    """

    from scipy.signal import welch

    seglen      = int(len(strain)*dt)
    srate       = int(1./dt)
    nperseg     = int(subseglen * srate)
    noverlap    = int(overlap_fraction*nperseg)
    window      = np.blackman(nperseg)

    freqs, psd = welch(strain,
                       fs              = 1./dt,
                       window          = window,
                       nperseg         = nperseg,
                       noverlap        = noverlap,
                       return_onesided = True,
                       scaling         = 'density')

    return freqs, psd

def read_asd_from_txt(path):
    return np.genfromtxt(path, usecols=[0,1], unpack=True)

def get_design_sensitivity(ifo):

    import os
    from ... import __path__
    main_path = os.path.abspath(__path__[0]+'/pipe/data/gw/asd/design/')

    if ifo=='L1' or ifo=='H1' or ifo=='I1':
        filename = 'LIGO-P1200087-v18-aLIGO_DESIGN.txt'
    elif ifo=='V1':
        filename = 'LIGO-P1200087-v18-AdV_DESIGN.txt'
    elif ifo=='K1':
        filename = 'LIGO-T1600593-v1-KAGRA_Design.txt'
    elif ifo=='ET':
        filename = 'LIGO-P1600143-v18-ET_D.txt'
    elif ifo=='CE':
        filename = 'LIGO-P1600143-v18-CE.txt'
    else:
        raise AttributeError("Design ASD not available for requested detector.\nDsign ASD is available for the following IFOs: H1, L1, V1, K1, ET, CE.")
    asd_path = os.path.join(main_path , filename)
    return np.genfromtxt(asd_path , usecols=[0,1], unpack=True)

def get_event_sensitivity(ifo, event):

    import os
    from ... import __path__
    asd_path        = os.path.abspath(__path__[0]+'/pipe/data/gw/asd/events/{}/{}_ASD.txt'.format(event, ifo))

    return np.genfromtxt(asd_path , usecols=[0,1], unpack=True)

def get_event_calibration(ifo, event):

    import os
    from ... import __path__
    spcal_path      = os.path.abspath(__path__[0]+'/pipe/data/gw/spcal/events/{}/{}_CalEnv.txt'.format(event, ifo))

    f, mag, phi, sigma_mag_low, sigma_phi_low, sigma_mag_up, sigma_phi_up  = np.transpose(np.genfromtxt(spcal_path , usecols=[0,1,2,3,4,5,6]))
    amp_sigma = (sigma_mag_up - sigma_mag_low)/2.
    phi_sigma = (sigma_phi_up - sigma_phi_low)/2.
    return [f,amp_sigma,phi_sigma]

class Noise(object):
    """
        Gaussian and stationary noise
    """

    def __init__(self, freqs, asd, f_min=None, f_max=None, filter=False):

        from scipy.interpolate import interp1d

        self.freqs          = freqs
        self.df             = np.median(np.diff(self.freqs))
        self.f_Nyq          = np.max(self.freqs)

        if f_max == None:
            self.f_max = self.f_Nyq
        else:
            self.f_max = f_max

        if f_min == None:
            self.f_min = np.min(self.freqs)
        else:
            self.f_min = f_min

        if filter == True:
            self.amp_spectrum   = filtering(self.freqs, asd, [self.f_min,self.f_max], type='bandpass', order=4)
        else:
            self.amp_spectrum   = asd

        self.power_spectrum = self.amp_spectrum*self.amp_spectrum

        # enlong ASD/PSD to f=0 to avoid problems during interpolation
        l_pad               = int(round(np.min(self.freqs)/self.df))+1
        self.freqs_pad      = np.append(np.flip(np.min(self.freqs) - np.arange(1,l_pad+1)*self.df), self.freqs)
        self.amp_spec_pad   = np.append(np.ones(l_pad)*self.amp_spectrum[0], self.amp_spectrum)
        self.pow_spec_pad   = np.append(np.ones(l_pad)*self.power_spectrum[0], self.power_spectrum)

        # check if input f_max is bigger than max(freq),
        # if yes, enlong ASD/PSD
        if self.f_max*2 > np.max(self.freqs_pad):
            l_pad               = int(round((self.f_max*2 - np.max(self.freqs_pad))/self.df))+1
            self.freqs_pad      = np.append(self.freqs_pad, np.max(self.freqs_pad)+np.arange(1,l_pad+1)*self.df)
            self.amp_spec_pad   = np.append(self.amp_spec_pad, np.ones(l_pad)*self.amp_spectrum[-1])
            self.pow_spec_pad   = np.append(self.pow_spec_pad,np.ones(l_pad)*self.power_spectrum[-1])

        # obs: linear interpolation is performed,
        # because higher orders give ASD/PSD < 0 (and this may screw up everything)
        self.psd_interp_func        = interp1d(self.freqs, self.power_spectrum)
        self.asd_interp_func        = interp1d(self.freqs, self.amp_spectrum)
        self.asd_interp_func_pad    = interp1d(self.freqs_pad, self.amp_spec_pad)
        self.psd_interp_func_pad    = interp1d(self.freqs_pad, self.pow_spec_pad)

    def interp_psd(self, fr):
        asd = self.asd_interp_func(fr)
        return asd * asd

    def interp_asd(self, fr):
        return self.asd_interp_func(fr)

    def interp_asd_pad(self, fr):
        return self.asd_interp_func_pad(fr)

    def interp_psd_pad(self, fr):
        return self.psd_interp_func_pad(fr)

    # to be updated, see generate_noise_gauss.py
    def generate_fake_noise(self, seglen, srate=4096., t_gps=1126259462., filter=False):
        """
            Generate fake gaussian noise from ASD.

            Arguments:
                - seglen  : period of the output segment
                - srate   : sampling rate of the output segment (optional, default 4096)
                - t_gps   : GPS time, center of the time series (optional, default 1126259462, i.e. GW150914)

            Return:
                - fake_series : numpy.array conatining the time series of the artificial noise strain

        """
        dt      = 1./srate
        f_max   = srate/2.
        df      = 1./seglen
        N_FD    = int(round(f_max/df))
        fr_out  = np.arange(N_FD+1)*df

        # filter PSD
        if filter:
            psd = filtering(fr_out, self.interp_psd_pad(fr_out), [self.f_min,self.f_max], type='bandpass', order=4)
        else:
            psd = self.interp_psd_pad(fr_out)

        # remove contributions below f_min
        psd[np.where(fr_out<self.f_min*0.95)] = 0.
        sigma       = 0.5*np.sqrt(psd/df)

        # generate (two) noise series
        noise1      = 1j * np.random.normal(np.zeros(len(sigma)), sigma)
        noise1      += np.random.normal(np.zeros(len(sigma))    , sigma)
        series1     = Series('freq' , noise1 , srate=srate, seglen=seglen ,
                             f_min=self.f_min, f_max=f_max, t_gps=t_gps, alpha_taper=0., filter=False)

        noise2      = 1j * np.random.normal(np.zeros(len(sigma)), sigma)
        noise2      += np.random.normal(np.zeros(len(sigma))    , sigma)
        series2     = Series('freq' , noise2 , srate=srate, seglen=seglen ,
                             f_min=self.f_min, f_max=f_max, t_gps=t_gps, alpha_taper=0., filter=False)

        # feather together noise1 & noise2
        # (in this way we avoid unwanted symmetries, specularities, repetitions, etc)
        x           = np.cos(np.pi*np.arange(len(series1.time_series))/len(series1.time_series))
        y           = np.sin(np.pi*np.arange(len(series2.time_series))/len(series2.time_series))
        noise       = x*series1.time_series + y*series2.time_series

        # create output series (without window, since it will be applied during the data reading)
        fake_series = Series('time' , noise , srate=srate, seglen=seglen , f_min=self.f_min, f_max=f_max, t_gps=t_gps, alpha_taper=0.)
        return fake_series.time_series

max_stars_repo_name: jacopok/bajes

max_issues_repo_name: jacopok/bajes

max_forks_repo_name: jacopok/bajes


text:
from typing import Optional, Dict, Type
from abc import ABC, abstractmethod
import platform
import time
import warnings
from enum import Enum

import numpy as np

from pyvirtualcam.util import FPSCounter, encode_fourcc, decode_fourcc

class Backend(ABC):
    """
    Describes the interface of Backend classes.

    Note that Backend classes should not inherit from this class,
    it exists only for documentation and static typing purposes.
    """

    @abstractmethod
    def __init__(self, *, width: int, height: int, fps: float,
                 fourcc: int, device: Optional[str], **kw):
        """
        All arguments are keyword-only arguments to avoid ambiguities.

        :param width: Image width of the frames passed to :meth:`send`.
            If the virtual camera does not support the given width and
            height combination then an exception must be raised.
        :param height: Image height of the frames passed to :meth:`send`.
        :param fps: Approximate rate at which :meth:`send` will be called,
            actual rate may be more or less. The backend should
            not expect a specific rate but may use this rate as a hint
            or an upper bound. If the rate is not one of the native frame
            rates supported by the virtual camera device, then a suitable
            rate should be chosen. An exception should generally not be raised.
        :param fourcc: Pixel format of frames passed to :meth:`send`, encoded as
            `libyuv FourCC <https://chromium.googlesource.com/libyuv/libyuv/+/refs/heads/master/include/libyuv/video_common.h>`_
            code. One of the codes in the :class:`~pyvirtualcam.PixelFormat` enum.
            If the pixel format is unsupported then an exception must
            be raised.
        :param device: If given, name of the virtual camera device to use,
            otherwise any available device can be used.
            Depending on the operating system, this is typically a device
            file name or a camera name as it appears in apps.
            If a device name is given then the specified device must be
            used or an exception raised if the device is unavailable.
            If no device name is given (``None``) and no device is available
            then an exception must be raised.
        :param kw: Extra keyword arguments passed through from user code.
        """

    @abstractmethod
    def close(self):
        """ Releases any resources associated with the virtual camera device.

        This method is guaranteed to be called exactly once.
        After this method was called, no further methods are called.
        """

    @abstractmethod
    def send(self, frame: np.ndarray):
        """ Send the given frame to the camera device.

        :param frame: A 1D C-contiguous uint8 numpy array corresponding
            to the chosen pixel format and frame width and height.
        """

    @abstractmethod
    def device(self) -> str:
        """ The name of the virtual camera device in use.

        If ``device`` was not ``None`` in the constructor, then
        the returned value must match the constructor argument.
        """

    @abstractmethod
    def native_fourcc(self) -> Optional[int]:
        """ The `libyuv FourCC <https://chromium.googlesource.com/libyuv/libyuv/+/refs/heads/master/include/libyuv/video_common.h>`_
        code of the native pixel format used in the backend.
        Must be one of the codes in the :class:`~pyvirtualcam.PixelFormat` enum,
        or ``None`` if no matching code is defined or the value
        for some other reason is not available or sensible.

        This does not necessarily correspond to all the formats
        that the device supports when a program captures from it,
        which is typical on Windows.
        Rather it is an indication on whether any pixel format
        conversions have to be done in the backend itself before
        sending the frame to the camera device.
        """

BACKENDS: Dict[str, Type[Backend]] = {}

def register_backend(name: str, clazz):
    """
    Register a new backend.

    If a backend with the same name is already registered,
    it will be replaced.

    :param name: Name of the backend.
        Used as ``backend`` argument in :meth:`~pyvirtualcam.Camera`.
    :param clazz: Class type of the backend conforming to :class:`~pyvirtualcam.Backend`.
    """
    BACKENDS[name] = clazz

if platform.system() == 'Windows':
    from pyvirtualcam import _native_windows_obs, _native_windows_unity_capture
    register_backend('obs', _native_windows_obs.Camera)
    register_backend('unitycapture', _native_windows_unity_capture.Camera)
elif platform.system() == 'Darwin':
    from pyvirtualcam import _native_macos_obs
    register_backend('obs', _native_macos_obs.Camera)
elif platform.system() == 'Linux':
    from pyvirtualcam import _native_linux_v4l2loopback
    register_backend('v4l2loopback', _native_linux_v4l2loopback.Camera)

class PixelFormat(Enum):
    """ Pixel formats.

    The enum values are libyuv FourCC codes.
    They are only used internally by the virtual camera backends.
    """

    # See external/libyuv/include/libyuv/video_common.h for fourcc codes.
    RGB = 'raw '
    """ Shape: ``(h,w,3)`` """

    BGR = '24BG'
    """ Shape: ``(h,w,3)`` """

    RGBA = 'ABGR'
    """ Shape: ``(h,w,4)`` """

    GRAY = 'J400'
    """ Shape: ``(h,w)`` """

    I420 = 'I420'
    """ Shape: any of size ``w * h * 3/2`` """

    NV12 = 'NV12'
    """ Shape: any of size ``w * h * 3/2`` """

    YUYV = 'YUY2'
    """ Shape: any of size ``w * h * 2`` """

    UYVY = 'UYVY'
    """ Shape: any of size ``w * h * 2`` """

    def __str__(self):
        return self.name

    def __repr__(self):
        return f'{type(self).__name__}.{self.name}'

FrameShapes = {
    PixelFormat.RGB: lambda w, h: (h, w, 3),
    PixelFormat.BGR: lambda w, h: (h, w, 3),
    PixelFormat.RGBA: lambda w, h: (h, w, 4),
    PixelFormat.GRAY: lambda w, h: (h, w),
    PixelFormat.I420: lambda w, h: w * h * 3 // 2,
    PixelFormat.NV12: lambda w, h: w * h * 3 // 2,
    PixelFormat.YUYV: lambda w, h: w * h * 2,
    PixelFormat.UYVY: lambda w, h: w * h * 2,
}

class Camera:
    """
    :param width: Frame width in pixels.
    :param height: Frame height in pixels.
    :param fps: Target frame rate in frames per second.
    :param fmt: Input pixel format.
    :param device: The virtual camera device to use.
        If ``None``, the first available device is used.

        Built-in backends:

        - ``v4l2loopback`` (Linux): ``/dev/video<n>``
        - ``obs`` (macOS/Windows): ``OBS Virtual Camera``
        - ``unitycapture`` (Windows): ``Unity Video Capture``, or the name you gave to the device
    :param backend: The virtual camera backend to use.
        If ``None``, all available backends are tried.

        Built-in backends:

        - ``v4l2loopback`` (Linux)
        - ``obs`` (macOS/Windows)
        - ``unitycapture`` (Windows)
    :param print_fps: Print frame rate every second.
    :param kw: Extra keyword arguments forwarded to the backend.
        Should only be given if a backend is specified.
        Note that the built-in backends do not have extra arguments.
    """
    def __init__(self, width: int, height: int, fps: float, *,
                 fmt: PixelFormat=PixelFormat.RGB,
                 device: Optional[str]=None,
                 backend: Optional[str]=None,
                 print_fps: bool=False,
                 **kw) -> None:
        if backend:
            backends = [(backend, BACKENDS[backend])]
        else:
            backends = list(BACKENDS.items())
        self._backend = None
        errors = []
        for name, clazz in backends:
            try:
                self._backend = clazz(
                    width=width, height=height, fps=fps,
                    fourcc=encode_fourcc(fmt.value),
                    device=device,
                    **kw)
            except Exception as e:
                errors.append(f"'{name}' backend: {e}")
            else:
                self._backend_name = name
                break
        if self._backend is None:
            raise RuntimeError('\n'.join(errors))

        self._width = width
        self._height = height
        self._fps = fps
        self._fmt = fmt
        self._print_fps = print_fps

        frame_shape = FrameShapes[fmt](width, height)
        if isinstance(frame_shape, int):
            def check_frame_shape(frame: np.ndarray):
                if frame.size != frame_shape:
                    raise ValueError(f"unexpected frame size: {frame.size} != {frame_shape}")
        else:
            def check_frame_shape(frame: np.ndarray):
                if frame.shape != frame_shape:
                    raise ValueError(f"unexpected frame shape: {frame.shape} != {frame_shape}")

        self._check_frame_shape = check_frame_shape

        self._fps_counter = FPSCounter(fps)
        self._fps_last_printed = time.perf_counter()
        self._frames_sent = 0
        self._last_frame_t = None
        self._extra_time_per_frame = 0

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> bool:
        self.close()
        return False

    def __del__(self):
        self.close()

    @property
    def backend(self) -> str:
        """ The virtual camera backend in use.
        """
        return self._backend_name

    @property
    def device(self) -> str:
        """ The virtual camera device in use.
        """
        return self._backend.device()

    @property
    def width(self) -> int:
        """ Frame width in pixels.
        """
        return self._width

    @property
    def height(self) -> int:
        """ Frame height in pixels.
        """
        return self._height

    @property
    def fps(self) -> float:
        """ Target frame rate in frames per second.
        """
        return self._fps

    @property
    def fmt(self) -> PixelFormat:
        """ The input pixel format as accepted by :meth:`send`.
        """
        return self._fmt

    @property
    def native_fmt(self) -> Optional[PixelFormat]:
        """ The native pixel format of the virtual camera,
        or ``None`` if not known or otherwise unavailable.

        Note that this is the native format in use in the virtual
        camera backend and may not necessarily be the native format
        of the virtual camera device itself.
        For example, on Windows, a camera device typically
        supports multiple formats.
        """
        fourcc = self._backend.native_fourcc()
        return PixelFormat(decode_fourcc(fourcc)) if fourcc else None

    @property
    def frames_sent(self) -> int:
        """ Number of frames sent.
        """
        return self._frames_sent

    def close(self) -> None:
        """ Manually free resources.

        This method is automatically called when using ``with`` or
        when this instance goes out of scope.
        """
        if self._backend is not None:
            self._backend.close()
            self._backend = None

    def send(self, frame: np.ndarray) -> None:
        """Send a frame to the virtual camera device.

        :param frame: Frame to send. The shape of the array must match
            the chosen :class:`~pyvirtualcam.PixelFormat`.
        """
        if frame.dtype != np.uint8:
            raise TypeError(f'unexpected frame dtype: {frame.dtype} != uint8')

        self._check_frame_shape(frame)

        self._frames_sent += 1
        self._last_frame_t = time.perf_counter()
        self._fps_counter.measure()

        if self._print_fps and self._last_frame_t - self._fps_last_printed > 1:
            self._fps_last_printed = self._last_frame_t
            s = f'{self._fps_counter.avg_fps:.1f} fps'

            # If sleep_until_next_frame() is used, show percentage of frame time
            # spent in computation (vs sleeping).
            if self._extra_time_per_frame > 0:
                busy_ratio = min(1, self._extra_time_per_frame * self._fps)
                s += f' | {100*busy_ratio:.0f} %'

            print(s)

        frame = np.array(frame.reshape(-1), copy=False, order='C')
        self._backend.send(frame)

    @property
    def current_fps(self) -> float:
        """ Current measured frames per second. """
        return self._fps_counter.avg_fps

    def sleep_until_next_frame(self) -> None:
        """ Adaptively sleep until the next frame is due.

        This method continuously adapts the sleep duration
        in order to reach the target frame rate.
        As a side effect, it estimates the time spent in computation
        which is printed as a percentage if ``print_fps=True``
        is given as argument in the constructor.
        """
        next_frame_t = self._last_frame_t + 1 / self._fps
        current_t = time.perf_counter()
        if current_t < next_frame_t:
            factor = self._fps / self._fps_counter.avg_fps - 1
            self._extra_time_per_frame += 0.01 * factor
            self._extra_time_per_frame = max(0, self._extra_time_per_frame)
            t_sleep = next_frame_t - current_t - self._extra_time_per_frame
            if t_sleep > 0:
                time.sleep(t_sleep)

ext: py

reason: import numpy


text:
import vtk, sys, os, numpy as np

import utilities
import extendField, pathlib, os, pandas, gdist
import matplotlib.pyplot as plt
import logging, argparse, pickle, scipy, scipy.stats
from vtk.util import numpy_support

def concatenateDict(*args):
    dRes = {}
    for d in args:
        dRes.update(d)
    return dRes

def set_consistent_normals(vtkMesh, splitting = False):
    """
    Checks that the orientation is consistent between all cells.
    """
    vtkNormals = vtk.vtkPolyDataNormals()
    vtkNormals.SetSplitting(splitting)
    vtkNormals.ConsistencyOn()
    vtkNormals.SetInputData(vtkMesh)
    vtkNormals.Update()
    return vtkNormals.GetOutput()

def compute_volume_vtk(vtkMesh, scale = 1e-3):

    massProps = vtk.vtkMassProperties()
    massProps.SetInputData(set_consistent_normals(vtkMesh))
    massProps.Update()
    return massProps.GetVolume() * scale

def meshPartitionGeodesics(mesh, writePath, pointsPulmonary, pointsTricuspid, apexId):
    """
    Creates a mesh in PLY (need for the tetrahedralization) and another one in VTK with the fields, because the PLY doesn't store the fields. Space is cheap, anyway.

    writePath needs to be the file path without extension.
    """

    points, faces = utilities.vtk_to_numpy(mesh, returnFaces= True, flatten= False)
    faces = faces.astype(np.int32)
    distancePulmonary = gdist.compute_gdist(points, faces, pointsPulmonary.astype(np.int32))
    distanceTricuspid = gdist.compute_gdist(points, faces, pointsTricuspid.astype(np.int32))
    distanceApex = gdist.compute_gdist(points, faces, np.array([apexId], dtype = np.int32))

    utilities.add_scalar(mesh,distancePulmonary, 'distancePulmonary')
    utilities.add_scalar(mesh,distanceApex, 'distanceApex')
    utilities.add_scalar(mesh,distanceTricuspid, 'distanceTricuspid')

    if writePath:
        utilities.write_poly(mesh,writePath + '.vtk', scalarFields = [distancePulmonary,  distanceApex, distanceTricuspid],
                             scalarFieldNames = ['distancePulmonary',  'distanceApex', 'distanceTricuspid']
                                 , format = 'vtk')
        utilities.write_poly(mesh,writePath, format = 'ply')

    return mesh

def propagateScalars(mesh, writePath, referenceMesh):
    arraysToCopy = {}
    arraysToCopy['distancePulmonary'] = numpy_support.vtk_to_numpy(referenceMesh.GetPointData().GetArray('distancePulmonary'))
    arraysToCopy['distanceApex'] = numpy_support.vtk_to_numpy(referenceMesh.GetPointData().GetArray('distanceApex'))
    arraysToCopy['distanceTricuspid'] = numpy_support.vtk_to_numpy(referenceMesh.GetPointData().GetArray('distanceTricuspid'))
    for k, c in arraysToCopy.items():
        utilities.add_scalar(mesh, c, k)
    if writePath:
        utilities.write_poly(mesh,writePath + '.vtk', scalarFields = list(arraysToCopy.values()),
                             scalarFieldNames = list(arraysToCopy.keys()), format = 'vtk')
        utilities.write_poly(mesh,writePath, format = 'ply')
    return mesh

def subVolumePartition(mesh,vtkName, writeFolder, rvotMethod = 'geodesic', **kwargs):
    """
    mesh: vtk mesh that will be partitioned
    vtkName: name of the vtk mesh (without the .vtk)

    Optional arguments, one of thw two needs to be present
    vtkMeshTemplate : use a template and point to point correspondence to propagate labels
    anatomicLabels: computes the partition using the landmarks
    """
    recompute = kwargs.get('recompute', False)
    geodesicPath = os.path.join(writeFolder, vtkName)
    if 'anatomicLabels' in kwargs:
        anatomicLabels = kwargs['anatomicLabels']
        meshPartitionGeodesics(mesh, geodesicPath,
                               anatomicLabels['pointsPulmonary'],  anatomicLabels['pointsTricuspid'],  anatomicLabels['apexId'])

    elif 'vtkMeshTemplate' in kwargs:
        mesh = propagateScalars(mesh, geodesicPath, kwargs['vtkMeshTemplate'])
    else:
        raise ValueError('No anatomic labels or template')
    vtkVolumetricPath = extendField.surfaceToVolumetric(mesh,
                                                        writeFolder, writeFolder, recompute = recompute)
    if not vtkVolumetricPath:
        return 0,0,0
    logging.info('outputPath = %s' % os.path.join(writeFolder, vtkName))
    vtkVolumetricWithFieldsPath = extendField.extendScalarField(vtkVolumetricPath, mesh, meshNameType = '_tetra',
                                  recompute = recompute)
    rvot, inlet, apex = extendField.computeGeodesicPartitionVolumes(vtkVolumetricWithFieldsPath)
    return rvot, inlet, apex

def readLabelsRVTOMTEC():
    """
    Read labels from hardcoded directories. Files needed:
    - Data/0_referenceADDETIA_CELLS.vtk
    - Data/pointsSeptum.csv
    - Data/facesTricuspid2.csv
    - Data/facesPulmonary2.csv
    """
    _, triangles = utilities.vtk_to_numpy(utilities.read_poly('Data/0_referenceADDETIA_CELLS.vtk'), returnFaces= True) # Just a mesh to get its topology
    facesTricuspid = pandas.read_csv('./Data/facesTricuspid2.csv')
    facesPulmonary = pandas.read_csv('./Data/facesPulmonary2.csv')
    pointsTricuspid = set()
    for _, f in facesTricuspid.iterrows():
        t = triangles[f.vtkOriginalCellIds]
        pointsTricuspid.add(t[0])
        pointsTricuspid.add(t[1])
        pointsTricuspid.add(t[2])
    pointsTricuspid = np.array(list(pointsTricuspid))

    pointsPulmonary = set()
    for _, f in facesPulmonary.iterrows():
        t = triangles[f.vtkOriginalCellIds]
        pointsPulmonary.add(t[0])
        pointsPulmonary.add(t[1])
        pointsPulmonary.add(t[2])
    pointsPulmonary = np.array(list(pointsPulmonary))

    facesValve = np.concatenate([facesPulmonary.vtkOriginalCellIds.values, facesTricuspid.vtkOriginalCellIds.values])
    apexId = 906
    midTricuspid = 102
    midPulmonary = 63

    septumPointsIDs = pandas.DataFrame.from_csv('Data/pointsSeptum.csv').values.reshape(-1)
    facesSeptum = []
    for i,t in enumerate(triangles):
        if all(map(lambda pointID: pointID in septumPointsIDs , t)):
            facesSeptum.append(i)
    facesLateral = [i for i in range(len(triangles)) if i not in facesValve and i not in facesSeptum]
    anatomicLabels = {}
    anatomicLabels['apexId'] = apexId
    anatomicLabels['facesLateral'] = facesLateral
    anatomicLabels['facesSeptum'] = facesSeptum
    anatomicLabels['pointsTricuspid'] = pointsTricuspid
    anatomicLabels['facesValve'] = facesValve
    anatomicLabels['pointsPulmonary'] = pointsPulmonary
    anatomicLabels['pointsValve'] =np.concatenate([pointsPulmonary, pointsTricuspid])

    return anatomicLabels

def regionalEjectionFraction(volumes):
    ed_frame = 0
    es_frame = np.argmin(np.sum(volumes, axis =1))
    return (volumes[ed_frame, :] - volumes[es_frame, :])/volumes[ed_frame, :]

def computeEDVEF(meshes):
    """
    Computes from a
    """
    anatomicLabels = readLabelsRVTOMTEC()
    if isinstance(meshes, list):
        meshes = {k:m for k, m in enumerate(meshes)}
    readIfString = lambda s: utilities.read_poly(s) if isinstance(s, str) else s
    meshes = {int(i): readIfString(m)  for i,m in meshes.items()}
    meshesList = [meshes[i] for i in range(len(meshes))]
    volumesTotal = np.array([ compute_volume_vtk(m) for m in meshesList])
    es = np.argmin(volumesTotal)
    ed = 0
    mesh_0 = meshPartitionGeodesics(meshesList[0] ,writePath = '/tmp/meshName',
                                    pointsPulmonary =anatomicLabels['pointsPulmonary'],
                                    pointsTricuspid =anatomicLabels['pointsTricuspid'],
                                    apexId = anatomicLabels['apexId'])

    volumesED = subVolumePartition(meshesList[ed], 'meshName', '/tmp', vtkMeshTemplate = mesh_0, recompute = True)
    volumesES = subVolumePartition(meshesList[es], 'meshName', '/tmp', vtkMeshTemplate = mesh_0, recompute = True)
    volumesED, volumesES = np.array(volumesED), np.array(volumesES)
    volumesEF  =100 *  (volumesED - volumesES)/volumesES
    return np.concatenate([volumesED, volumesEF])

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Computes the regional volumes dynamics.')
    parser.add_argument('inputPath',
                        help=' input path with the vtk meshes')
    parser.add_argument('outputPathResult',  help='path to write the dictionary')

    parser.add_argument('outputPathMeshes',
                        help='path to write all the meshes')

    parser.add_argument('-rvotMethod', default = 'geodesic',
                        help='geodesic/heat')
    parser.add_argument('--recompute', action = 'store_true',
                        help='path to write all the meshes')

    args = parser.parse_args()

    #--------------------
    # Read the klabels
    #--------------------
    anatomicLabels = readLabelsRVTOMTEC()

    allVolumes = {}
    print('Recompute = ', args.recompute)
    for pId in os.listdir(args.inputPath):
        if '00.vtk' not in pId:
            continue
        pId = '_'.join(pId.split('_')[:-1]) + '_'
        print('pId = %s' % pId)
        meshes = utilities.read_meshes(args.inputPath,pId, convertToNumpy= False)
        meshes = {int(i): m  for i,m in meshes.items()}
        meshesList = [meshes[i] for i in range(len(meshes))]
        volumes = []
        for i, m in enumerate(meshesList):
            if i == 0:
                mesh_0 = meshPartitionGeodesics(m ,writePath = None,
                                                pointsPulmonary =anatomicLabels['pointsPulmonary'],
                                                pointsTricuspid =anatomicLabels['pointsTricuspid'],
                                                apexId = anatomicLabels['apexId'])

            volumes.append(subVolumePartition(meshesList[i], pId + '_%d' % i, args.outputPathMeshes, rvotMethod = args.rvotMethod,  vtkMeshTemplate = mesh_0, recompute = args.recompute))
        allVolumes[pId] = np.array(volumes)

    with open(args.outputPathResult, 'wb') as file:
        pickle.dump(allVolumes, file)

max_stars_repo_name: gbernardino/RVparcellation

max_issues_repo_name: gbernardino/RVparcellation

max_forks_repo_name: gbernardino/RVparcellation


text:
import numpy as np
import matplotlib.pyplot as plt

# We Will Solve the Diff Equation: dA/dt = -λA = f(A,t), 0<=t<=8 [s]
def f(t, A):
    λ = np.log(2)/2 # Constant
    f = -λ*A
    return f

# Second Order Runge Kutta Method (τ: the step)
def RK2(τ):
    # Interval Edges
    a, b = 0, 8
    # Number of Repetitions
    N = int((b-a)/τ)

    w = np.zeros(N+1)
    t = np.zeros(N+1)

    # Starting Points
    w[0], t[0] = 1, 0

    for i in range(N):
        w[i+1] = w[i] + τ * f(t[i] + τ/2, w[i] + τ/2*f(t[i],w[i]))
        t[i+1] = t[i] + τ

    return w, t

# We Compute the Uncertenty
def unc_RK2(τ):
    unc = np.zeros(11)
    rk2_τ, rk2_2τ = RK2(τ)[0], RK2(τ/2)[0]

    for i in range(1,11):
        unc[i] = (rk2_τ[i]-rk2_2τ[2*i])/3

    return unc

if __name__ == '__main__':

    w_rk2, t_rk2, unc_rk2 = RK2(0.8)[0], RK2(0.8)[1], unc_RK2(0.8)

    print('\nComputed with 2^nd Order Runge Kutta Method:')
    for i in range(11):
        print(f"\tStep {i} with t_{i}={t_rk2[i]:.1f}: {w_rk2[i]:.8f} ± {unc_rk2[i]:.8f}")

    # Visualize the Results
    ## Plot the Solutions for Both Methods
    plt.plot(t_rk2, w_rk2, color='black', marker=".", linewidth=1, label='Runge Kutta 2th Order')

    ## How Graph is Shown (legend, axis titles, graph title)
    plt.legend(loc='upper center', prop={'size': 10})
    plt.ylabel('A(t) [kBq]')
    plt.xlabel('Time [s]')
    plt.title('Solution of Differential Equation')

    ## Show the major grid lines with dark grey lines
    plt.grid(b=True, which='major', color='#666666', linestyle='--')

    ## Show the minor grid lines
    plt.minorticks_on()
    plt.grid(b=True, which='minor', color='#999999', linestyle='--', alpha=0.2)

    ## Fix Quality
    plt.tight_layout()

    ## Show the Graph
    plt.show()

max_stars_repo_name: AlexTsagas/Computational-Physics

max_issues_repo_name: AlexTsagas/Computational-Physics

max_forks_repo_name: AlexTsagas/Computational-Physics


text:
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

def data2feature(f_name,cla):
	file_value = f_name.values
	file_value[:,-1] = cla
	feature = file_value
	feature = feature[:,1:]
	np.random.shuffle(feature)
	return feature

f0 = pd.read_csv("./labeld_Monday-Benign.csv")
f1 = pd.read_csv("./labeld_Wednesday-Dos-GlodenEye-7350.csv")
f2 = pd.read_csv("./labeld_Wednesday-Dos-Hulk5919.csv")
f3 = pd.read_csv("./labeld_Friday-DDoS-16050.csv")
f4 = pd.read_csv("./labeld_Tuesday-FTP-Palator-3945.csv")
f5 = pd.read_csv("./labeld_Wednesday-Dos-slowloris-1916.csv")
f6 = pd.read_csv("./labeld_Wednesday-Dos-slowlhttptest-1403.csv")

type_0 = data2feature(f0,0)[:4]
type_1 = data2feature(f1,1)[:4]
type_2 = data2feature(f2,2)[:4]
type_3 = data2feature(f3,3)[:4]
type_4 = data2feature(f4,4)[:4]
type_5 = data2feature(f5,5)[:4]
type_6 = data2feature(f6,6)[:4]

imgs = np.concatenate((type_0,type_1,type_2,type_3,type_4,type_5,type_6),axis=0)

fig = plt.figure()
for i in range(7):
	for j in range(4):
		ax = fig.add_subplot(7,4,i*4+(j+1))
		cla = (imgs[i*4+j])[:-1]
		img_title = (imgs[i*4+j])[-1]
		img = cla.reshape([40,40])
		ax.imshow(img,cmap='gray')
		plt.axis("off")
		plt.title(img_title)

plt.show()


text:
import numpy as np
import pickle
from sklearn.utils import class_weight, shuffle
from sklearn.metrics import classification_report
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.linear_model import LogisticRegression
#import tensorflowjs as tfjs

# Seeding for reproducibility
np.random.seed(0)

# Update to the directory where you saved the
# preprocessed data
dataset_directory = "formated_data/"

# Reading data
reddit_file0 = open(dataset_directory + "reddit0.pickle", "rb")
reddit_data0 = pickle.load(reddit_file0)
reddit_file0.close()
reddit_file1 = open(dataset_directory + "reddit1.pickle", "rb")
reddit_data1 = pickle.load(reddit_file1)
reddit_file1.close()
reddit_file2 = open(dataset_directory + "reddit2.pickle", "rb")
reddit_data2 = pickle.load(reddit_file2)
reddit_file2.close()
reddit_file3 = open(dataset_directory + "reddit3.pickle", "rb")
reddit_data3 = pickle.load(reddit_file3)
reddit_file3.close()

reddit_data = reddit_data0 + reddit_data1 + reddit_data2 + reddit_data3 # Combine lists

hackernews_file = open(dataset_directory + "hacker_news.pickle", "rb")
hackernews_data = pickle.load(hackernews_file)
hackernews_file.close()

youtube_file = open(dataset_directory + "youtube.pickle", "rb")
youtube_data = pickle.load(youtube_file)
youtube_file.close()

reddit_samples = len(reddit_data)
hackernews_samples = len(hackernews_data)
youtube_samples = len(youtube_data)
samples = reddit_samples + hackernews_samples + youtube_samples

# Creating labels
reddit_labels = [0 for i in range(reddit_samples)]
hackernews_labels = [1 for i in range(hackernews_samples)]
youtube_labels = [2 for i in range(youtube_samples)]

# Concatinating data
data = reddit_data + hackernews_data + youtube_data
labels = reddit_labels + hackernews_labels + youtube_labels

# Shuffling data so that the distribution of train and test
# data are as similar as possible
data, labels = shuffle(data, labels, random_state=0)

# Splitting the data into train, validation and test.
x_train = data[0 : samples * 70 // 100]
x_val = data[samples * 70  // 100 : samples * 85 // 100]
x_test = data[samples * 85 // 100 :]

y_train = labels[0 : samples * 70 // 100]
y_val = labels[samples * 70  // 100 : samples * 85 // 100]
y_test = labels[samples * 85 // 100 :]

# Class weights so that smaller classes are given more weight per sample
class_weights = class_weight.compute_class_weight("balanced", [0, 1, 2], y=y_train)
class_weights = dict(enumerate(class_weights))

# Vectorize using the number of words of each type
bow_converter = CountVectorizer(tokenizer=lambda doc: doc)

x_train = bow_converter.fit_transform(x_train)
x_val = bow_converter.transform(x_val)
x_test = bow_converter.transform(x_test)

# Create model
model = LogisticRegression(class_weight=class_weights, max_iter=10000).fit(x_train, y_train)

# Print results
val_score = model.score(x_val, y_val)
print("Validation: [accuracy: " + str(val_score) + "]")
test_score = model.score(x_test, y_test)
print("Test: [accuracy: " + str(test_score) + "]")


text:
"""Module with utils for logging."""

import logging

import IPython.display
import ipywidgets as widgets
import numpy as np

def setup_log_print_options():
    """Set up print options for logging output."""
    np.set_printoptions(precision=3, suppress=True)

# Logging widgets
class OutputWidgetHandler(logging.Handler):
    """Logging handler sending logs to an output widget.

    This version is copied (with minor adaptations) from:
    https://github.com/ai4er-cdt/gtc-biodiversity/blob/main/geograph/visualisation/widget_utils.py
    Originally taken from:
    https://ipywidgets.readthedocs.io/en/latest/examples/Output%20Widget.html
    """

    def __init__(self, *args, max_len=30, **kwargs):
        """Logging handler sending logs to an output widget.

        Args:
            max_len (int, optional): maximum length of log in lines. Defaults to 30.
        """
        super().__init__(*args, **kwargs)
        self.max_len = max_len
        layout = {
            "width": "100%",
            "border": "1px solid black",
        }
        self.out = widgets.Output(layout=layout)
        formatter = logging.Formatter(
            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
        )
        self.setFormatter(formatter)

    def emit(self, record):
        """Overload of logging.Handler method."""
        formatted_record = self.format(record)
        new_output = {
            "name": "stdout",
            "output_type": "stream",
            "text": formatted_record + "\n",
        }
        if len(self.out.outputs) > self.max_len:
            self.out.outputs = self.out.outputs[1:] + (new_output,)
        else:
            self.out.outputs = self.out.outputs + (new_output,)

    def show_logs(self):
        """Show the logs."""
        IPython.display.display(self.out)

    def clear_logs(self):
        """Clear the current logs."""
        self.out.clear_output()

max_stars_repo_name: rdnfn/solar-agent

max_issues_repo_name: rdnfn/solar-agent

max_forks_repo_name: rdnfn/solar-agent


text:
import numpy as np

from opytimizer.optimizers.boolean import umda
from opytimizer.spaces import boolean

def test_umda_params():
    params = {
        'p_selection': 0.75,
        'lower_bound': 0.05,
        'upper_bound': 0.95
    }

    new_umda = umda.UMDA(params=params)

    assert new_umda.p_selection == 0.75

    assert new_umda.lower_bound == 0.05

    assert new_umda.upper_bound == 0.95

def test_umda_params_setter():
    new_umda = umda.UMDA()

    try:
        new_umda.p_selection = 'a'
    except:
        new_umda.p_selection = 0.75

    assert new_umda.p_selection == 0.75

    try:
        new_umda.p_selection = -1
    except:
        new_umda.p_selection = 0.75

    assert new_umda.p_selection == 0.75

    try:
        new_umda.lower_bound = 'a'
    except:
        new_umda.lower_bound = 0.05

    assert new_umda.lower_bound == 0.05

    try:
        new_umda.lower_bound = -1
    except:
        new_umda.lower_bound = 0.05

    assert new_umda.lower_bound == 0.05

    try:
        new_umda.upper_bound = 'a'
    except:
        new_umda.upper_bound = 0.95

    assert new_umda.upper_bound == 0.95

    try:
        new_umda.upper_bound = -1
    except:
        new_umda.upper_bound = 0.95

    assert new_umda.upper_bound == 0.95

    try:
        new_umda.upper_bound = 0.04
    except:
        new_umda.upper_bound = 0.95

    assert new_umda.upper_bound == 0.95

def test_umda_calculate_probability():
    new_umda = umda.UMDA()

    boolean_space = boolean.BooleanSpace(n_agents=5, n_variables=2)

    probs = new_umda._calculate_probability(boolean_space.agents)

    assert probs.shape == (2, 1)

def test_umda_sample_position():
    new_umda = umda.UMDA()

    probs = np.zeros((1, 1))

    position = new_umda._sample_position(probs)

    assert position == 1

def test_umda_update():
    new_umda = umda.UMDA()

    boolean_space = boolean.BooleanSpace(n_agents=2, n_variables=5)

    new_umda.update(boolean_space)

max_stars_repo_name: anukaal/opytimizer

max_issues_repo_name: anukaal/opytimizer

max_forks_repo_name: anukaal/opytimizer


text:
#======================================================================
#
#     This module contains routines to postprocess the VFI
#     solutions.
#
#     Simon Scheidegger, 01/19
#     edited by Patrick O'Callaghan, with Cameron Gordon and Josh Aberdeen, 11/2021
#======================================================================

import numpy as np
from parameters import *
#import cPickle as pickle
import pickle
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import (RBF, Matern, RationalQuadratic,
                                              ExpSineSquared, DotProduct,
                                              ConstantKernel)

from datetime import datetime
from numpy.random import PCG64

#======================================================================
# Routine compute the errors
def ls_error(n_agents, t1, t2, num_points):
    file=open('errors.txt', 'w')
    now = datetime.now()
    dt = int(now.strftime("%H%M%S%f"))
    print("Time seed = ", dt)
    rng = np.random.default_rng(dt)
    unif=rng.uniform(0, 1, (num_points, n_agents))
    #sample of states
    kap_smp = kap_L+(unif)*(kap_U-kap_L)
    to_print=np.empty((1,3))
    if (t1 == 1):
        t1+1
    for i in range(t1, t2-1):
        sum_diffs=0
        diff = 0
        unif=rng.uniform(0, 1, (num_points, n_agents))
        #sample of states
        kap_smp = kap_L+(unif)*(kap_U-kap_L)
        # Load the model from the previous iteration step
        restart_data = filename + str(i) + ".pcl"
        with open(restart_data, 'rb') as fd_old:
            gp_old = pickle.load(fd_old)
            print("data from iteration step ", i , "loaded from disk")
        fd_old.close()

        # Load the model from the previous iteration step
        restart_data = filename + str(i+1) + ".pcl"
        with open(restart_data, 'rb') as fd_new:
            gp_new = pickle.load(fd_new)
            print("data from iteration step ", i+1 , "loaded from disk")
        fd_new.close()

        y_pred_old, sigma_old = gp_old.predict(kap_smp, return_std=True)
        y_pred_new, sigma_new = gp_new.predict(kap_smp, return_std=True)

        # plot predictive mean and 95% quantiles
        #for j in range(num_points):
            #print kap_smp[j], " ",y_pred_new[j], " ",y_pred_new[j] + 1.96*sigma_new[j]," ",y_pred_new[j] - 1.96*sigma_new[j]

        diff = y_pred_old-y_pred_new
        max_abs_diff=np.amax(np.fabs(diff))
        average = np.average(np.fabs(diff))

        to_print[0,0]= i+1
        to_print[0,1]= max_abs_diff
        to_print[0,2]= average

        np.savetxt(file, to_print, fmt='%2.16f')
        np.set_printoptions(suppress=True)
        msg=str(diff) + ", max=" + str(max_abs_diff)
        print(msg)
        print("===================================")

    file.close()

    return

#======================================================================

max_stars_repo_name: uq-aibe/aim

max_issues_repo_name: uq-aibe/aim

max_forks_repo_name: uq-aibe/aim

reason: import numpy,from numpy


text:
# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
from ..utils.scripts import get_parser

__all__ = ['derived_images']

def main(args=None):
    parser = get_parser(derived_images)
    parser.add_argument('infile', type=str,
                        help='Input FITS file name')
    parser.add_argument('outfile', type=str,
                        help='Output FITS file name')
    parser.add_argument('theta', type=float,
                        help='On-region correlation radius (deg)')
    parser.add_argument('--is_off_correlated', action='store_false',
                        help='Are the basic OFF maps correlated?')
    parser.add_argument('--overwrite', action='store_true',
                        help='Overwrite existing output file?')
    args = parser.parse_args(args)
    derived_images(**vars(args))

def derived_images(infile,
                   outfile,
                   theta,
                   is_off_correlated,
                   overwrite):
    """Make derived maps for a given set of basic maps.

    TODO: describe
    """
    import logging
    logging.basicConfig(level=logging.DEBUG, format='%(levelname)s - %(message)s')
    from astropy.io import fits
    from gammapy.background import Maps

    logging.info('Reading {0}'.format(infile))
    hdus = fits.open(infile)
    maps = Maps(hdus, theta=theta,
                is_off_correlated=is_off_correlated)
    logging.info('Computing derived maps')
    maps.make_derived_maps()
    logging.info('Writing {0}'.format(outfile))
    maps.writeto(outfile, clobber=overwrite)

max_stars_repo_name: joleroi/gammapy

max_issues_repo_name: joleroi/gammapy

max_forks_repo_name: joleroi/gammapy

reason: from astropy


text:
import numpy as np
import cv2 as cv

img = cv.imread('1.jpeg', cv.IMREAD_COLOR)
cv.line(img, (0,34), (44,65), (255,255,255), 8)

cv.imshow('image',img)
cv.waitKey(0)
cv.destroyAllWindows()


text:
import keras
from keras.models import Sequential, Model, load_model
from keras import backend as K

import tensorflow as tf

import os

import isolearn.keras as iso
import numpy as np

def logit(x) :
	return np.log(x / (1.0 - x))

def get_aparent_encoder() :
	onehot_encoder = iso.OneHotEncoder(205)

	def encode_for_aparent(sequences) :
		one_hots = np.concatenate([np.reshape(onehot_encoder(sequence), (1, len(sequence), 4, 1)) for sequence in sequences], axis=0)

		return [
			one_hots,
			np.zeros((len(sequences), 13)),
			np.ones((len(sequences), 1))
		]

	return encode_for_aparent

def get_aparent_legacy_encoder() :
	onehot_encoder = iso.OneHotEncoder(185)

	def encode_for_aparent(sequences) :
		one_hots = np.concatenate([np.reshape(onehot_encoder(sequence), (1, 1, len(sequence), 4)) for sequence in sequences], axis=0)

		return [
			one_hots,
			np.zeros((len(sequences), 36)),
			np.ones((len(sequences), 1))
		]

	return encode_for_aparent

def get_apadb_encoder() :
	onehot_encoder = iso.OneHotEncoder(205)

	def encode_for_apadb(prox_sequences, dist_sequences, prox_cut_starts, prox_cut_ends, dist_cut_starts, dist_cut_ends, site_distances) :
		prox_one_hots = np.concatenate([np.reshape(onehot_encoder(sequence), (1, len(sequence), 4, 1)) for sequence in prox_sequences], axis=0)
		dist_one_hots = np.concatenate([np.reshape(onehot_encoder(sequence), (1, len(sequence), 4, 1)) for sequence in dist_sequences], axis=0)

		return [
			prox_one_hots,
			dist_one_hots,
			np.array(prox_cut_starts).reshape(-1, 1),
			np.array(prox_cut_ends).reshape(-1, 1),
			np.array(dist_cut_starts).reshape(-1, 1),
			np.array(dist_cut_ends).reshape(-1, 1),
			np.log(np.array(site_distances).reshape(-1, 1)),
			np.zeros((len(prox_sequences), 13)),
			np.ones((len(prox_sequences), 1))
		]

	return encode_for_apadb


text:
############################TESTS ON POTENTIALS################################
from __future__ import print_function, division
import os
import sys
import numpy
from galpy import potential
from galpy import df
from galpy.potential import SCFPotential
from galpy.util import coords
from galpy.orbit import Orbit
_TRAVIS= bool(os.getenv('TRAVIS'))

EPS = 1e-13 ## default epsilon

DEFAULT_R= numpy.array([0.5,1.,2.])
DEFAULT_Z= numpy.array([0.,.125,-.125,0.25,-0.25])
DEFAULT_PHI= numpy.array([0.,0.5,-0.5,1.,-1.,
                       numpy.pi,0.5+numpy.pi,
                       1.+numpy.pi])

##Tests whether invalid coefficients will throw an error at runtime

def test_coeffs_toomanydimensions():
    Acos = numpy.ones((10,2,32,34))
    try:
        SCFPotential(Acos=Acos)
        raise Exception("Expected RuntimeError")
    except RuntimeError:
        pass

def test_coeffs_toolittledimensions():
    Acos = numpy.ones((10,2))
    try:
        SCFPotential(Acos=Acos)
        raise Exception("Expected RuntimeError")
    except RuntimeError:
        pass

def test_coeffs_AsinNone_LnotequalM():
    Acos = numpy.ones((2,3,4))
    try:
        SCFPotential(Acos=Acos)
        raise Exception("Expected RuntimeError")
    except RuntimeError:
        pass

def test_coeffs_AsinNotNone_LnotequalM():
    Acos = numpy.ones((2,3,4))
    Asin = numpy.ones((2,3,4))

    try:
        SCFPotential(Acos=Acos, Asin=Asin)
        raise Exception("Expected RuntimeError")
    except RuntimeError:
        pass

def test_coeffs_AsinNone_Mequals1():
    Acos = numpy.zeros((2,3,1))
    Asin = None

    SCFPotential(Acos=Acos, Asin=Asin)

def test_coeffs_AsinNone_MequalsL():
    Acos = numpy.zeros((2,3,3))
    Asin = None

    SCFPotential(Acos=Acos, Asin=Asin)

def test_coeffs_AsinNone_AcosNotaxisym():
    Acos = numpy.ones((2,3,3))
    try:
        SCFPotential(Acos=Acos)
        raise Exception("Expected RuntimeError")
    except RuntimeError:
        pass

def test_coeffs_AsinShape_notequal_AcosShape():
    Acos = numpy.ones((2,3,3))
    Asin = numpy.ones((2,2,2))
    try:
        SCFPotential(Acos=Acos, Asin=Asin)
        raise Exception("Expected RuntimeError")
    except RuntimeError:
        pass

def test_coeffs_Acos_L_M_notLowerTriangular():
    Acos = numpy.ones((2,3,3))
    Asin = numpy.zeros((2,3,3))
    try:
        SCFPotential(Acos=Acos, Asin=Asin)
        raise Exception("Expected RuntimeWarning")
    except RuntimeWarning:
        pass

def test_coeffs_Asin_L_M_notLowerTriangular():
    Acos = numpy.zeros((2,3,3))
    Asin = numpy.ones((2,3,3))
    try:
        SCFPotential(Acos=Acos, Asin=Asin)
        raise Exception("Expected RuntimeWarning")
    except RuntimeWarning:
        pass

def testAxi_phiIsNone():
    R = 1; z = 0; phi = 1.1;
    scf = SCFPotential()
    assert scf(R,z,None) == scf(R,z,phi), "The axisymmetric potential does not work at phi=None"
    assert scf.dens(R,z,None) == scf.dens(R,z,phi), "The axisymmetric density does not work at phi=None"
    assert scf.Rforce(R,z,None) == scf.Rforce(R,z,phi), "The axisymmetric Rforce does not work at phi=None"
    assert scf.zforce(R,z,None) == scf.zforce(R,z,phi), "The axisymmetric zforce does not work at phi=None"
    assert scf.phiforce(R,z,None) == scf.phiforce(R,z,phi), "The axisymmetric phiforce does not work at phi=None"

##Tests user inputs as arrays

def testArray_RArray():
    scf = SCFPotential()
    array = numpy.linspace(0, 3, 100)
    ArrayTest(scf, [array, 1., 0])

def testArray_zArray():
    scf = SCFPotential()
    array = numpy.linspace(0, 3, 100)
    ArrayTest(scf, [1., array, 0])

def testArray_phiArray():
    scf = SCFPotential()
    array = numpy.linspace(0, 3, 100)
    ArrayTest(scf, [1., 1., array])

def testArrayBroadcasting():
    scf = SCFPotential()
    R = numpy.ones((10,20,2))
    z = numpy.linspace(0,numpy.pi, 10)[:,None,None]
    phi = numpy.zeros((10,20))[:,:,None]

    ArrayTest(scf, [R, z, phi])

## tests whether scf_compute_spherical computes the correct coefficients for a Hernquist Potential
def test_scf_compute_spherical_hernquist():
    Acos, Asin = potential.scf_compute_coeffs_spherical(sphericalHernquistDensity, 10)
    spherical_coeffsTest(Acos, Asin)
    assert numpy.fabs(Acos[0,0,0] - 1.) < EPS, "Acos(n=0,l=0,m=0) = 1 fails. Found to be Acos(n=0,l=0,m=0) = {0}".format(Acos[0,0,0])
    assert numpy.all(numpy.fabs(Acos[1:,0,0]) < EPS), "Acos(n>0,l=0,m=0) = 0 fails."

## tests whether scf_compute_spherical computes the correct coefficients for Zeeuw's Potential
def test_scf_compute_spherical_zeeuw():
    Acos, Asin = potential.scf_compute_coeffs_spherical(rho_Zeeuw, 10)
    spherical_coeffsTest(Acos, Asin)
    assert numpy.fabs(Acos[0,0,0] - 2*3./4) < EPS, "Acos(n=0,l=0,m=0) = 3/2 fails. Found to be Acos(n=0,l=0,m=0) = {0}".format(Acos[0,0,0])
    assert numpy.fabs(Acos[1,0,0] - 2*1./12) < EPS, "Acos(n=1,l=0,m=0) = 1/6 fails. Found to be Acos(n=0,l=0,m=0) = {0}".format(Acos[0,0,0])
    assert numpy.all(numpy.fabs(Acos[2:,0,0]) < EPS), "Acos(n>1,l=0,m=0) = 0 fails."

##Tests that the numerically calculated results from axi_density1 matches with the analytic results
def test_scf_compute_axi_density1():
    A = potential.scf_compute_coeffs_axi(axi_density1, 10,10)
    axi_coeffsTest(A[0], A[1])
    analytically_calculated = numpy.array([[4./3, 7.* 3**(-5/2.), 2*11*5**(-5./2), 0],
                                            [0, 0,0,0],
                                            [0,11. / (3.**(5./2) * 5 * 7. * 2), 1. / (2*3.*5**.5*7.), 0]])
    numerically_calculated = A[0][:3,:4,0]
    shape = numerically_calculated.shape
    for n in range(shape[0]):
        for l in range(shape[1]):
            assert numpy.fabs(numerically_calculated[n,l] - analytically_calculated[n,l]) < EPS, \
        "Acos(n={0},l={1},0) = {2}, whereas it was analytically calculated to be {3}".format(n,l, numerically_calculated[n,l], analytically_calculated[n,l])
    #Checks that A at l != 0,1,2 are always zero
    assert numpy.all(numpy.fabs(A[0][:,3:,0]) < 1e-10), "Acos(n,l>2,m=0) = 0 fails."

    #Checks that A at n odd is always zero
    assert numpy.all(numpy.fabs(A[0][1::2,:,0]) < 1e-10), "Acos(n odd,l,m=0) = 0 fails."

    #Checks that A = 0 when n != 0 and l = 0
    assert numpy.all(numpy.fabs(A[0][1:,0,0]) < 1e-10), "Acos(n > 1,l=0,m=0) = 0 fails."

##Tests that the numerically calculated results from axi_density2 matches with the analytic results
def test_scf_compute_axi_density2():
    A = potential.scf_compute_coeffs_axi(axi_density2, 10,10,
                                         radial_order=30,costheta_order=12)
    axi_coeffsTest(A[0], A[1])
    analytically_calculated = 2*numpy.array([[1., 7.* 3**(-3/2.) /4., 3*11*5**(-5./2)/2., 0],
                                            [0,0,0,0], ##I never did analytically solve for n=1
                                            [0, 11./(7*5*3**(3./2)*2**(3.)), (7 * 5**(.5)*2**3.)**-1., 0]])
    numerically_calculated = A[0][:3,:4,0]
    shape = numerically_calculated.shape
    for n in range(shape[0]):
        if n ==1: continue
        for l in range(shape[1]):
            assert numpy.fabs(numerically_calculated[n,l] - analytically_calculated[n,l]) < EPS, \
        "Acos(n={0},l={1},0) = {2}, whereas it was analytically calculated to be {3}".format(n,l, numerically_calculated[n,l], analytically_calculated[n,l])

    #Checks that A at l != 0,1,2 are always zero
    assert numpy.all(numpy.fabs(A[0][:,3:,0]) < 1e-10), "Acos(n,l>2,m=0) = 0 fails."

    #Checks that A = 0 when n = 2,4,..,2*n and l = 0
    assert numpy.all(numpy.fabs(A[0][2::2,0,0]) < 1e-10), "Acos(n > 1,l = 0,m=0) = 0 fails."

## Tests how nbody calculation compares to density calculation for scf_compute_coeff_spherical
def compare_dens_nbody():
    N= int(1e5)
    Mh= 11.
    ah= 50./8.
    m= Mh/N
    factor=1.
    nsamp=10
    Norder=10
    Lorder=10

    hern= potential.HernquistPotential(amp=2*Mh,a=ah)
    hern.turn_physical_off()
    hdf= df.isotropicHernquistdf(hern)
    samp= [hdf.sample(n=N) for i in range(nsamp)]

    positions= numpy.array([[samp[i].x(),samp[i].y(),samp[i].z()*factor] for i in range(nsamp)])

    tptp= potential.TwoPowerTriaxialPotential(amp=2.*Mh,a=ah,alpha=1.,beta=4.,b=1.,c=factor)
    tptp.turn_physical_off()

    cc, ss= potential.scf_compute_coeffs(tptp.dens,Norder,Lorder,a=ah)
    c,s= numpy.zeros((2, nsamp, Norder, Lorder, Lorder))
    for i,p in enumerate(positions):
        c[i],s[i]= potential.scf_compute_coeffs_nbody(p,m*numpy.ones(N)*factor,Norder,Lorder,a=ah)

    # Check that the difference between the coefficients is within two standard deviations
    assert (cc-(numpy.mean(c,axis=0))<=(2.*numpy.std(c,axis=0))).all()

## Tests how nbody calculation compares to density calculation for scf_compute_coeff
def compare_dens_nbody_spherical():
    N= int(1e6)
    Mh= 11.
    ah= 50./8.
    m= Mh/N
    factor=1.
    nsamp=10
    Norder=10

    hern= potential.HernquistPotential(amp=2*Mh,a=ah)
    hern.turn_physical_off()
    hdf= df.isotropicHernquistdf(hern)
    samples= [hdf.sample(n=N) for i in range(nsamp)]

    positions= numpy.array([[samples[i].x(),samples[i].y(),samples[i].z()*factor] for i in range(nsamp)])

    c= numpy.zeros((nsamp,Norder,1,1))
    s= numpy.zeros((nsamp,Norder,1,1))
    for i in range(nsamp):
        c[i],s[i]= potential.scf_compute_coeffs_spherical_nbody(m=m*numpy.ones(N),pos=positions[i],N=Norder,a=ah)

    cc,ss= potential.scf_compute_coeffs_spherical(N=Norder,a=ah,dens=hern.dens)

    # Check that the difference between the coefficients is within the standard deviation
    assert (cc-numpy.mean(c,axis=0)<numpy.std(c,axis=0)).all()

def test_scf_compute_nfw():
    Acos, Asin = potential.scf_compute_coeffs_spherical(rho_NFW, 10)
    spherical_coeffsTest(Acos, Asin)

##Tests radial order from scf_compute_coeffs_spherical
def test_nfw_sphericalOrder():
    Acos, Asin = potential.scf_compute_coeffs_spherical(rho_NFW, 10)
    Acos2, Asin2 = potential.scf_compute_coeffs_spherical(rho_NFW, 10, radial_order=50)

    assert numpy.all(numpy.fabs(Acos - Acos2) < EPS), \
    "Increasing the radial order fails for scf_compute_coeffs_spherical"

##Tests radial and costheta order from scf_compute_coeffs_axi
def test_axi_density1_axiOrder():
    Acos, Asin = potential.scf_compute_coeffs_axi(axi_density1, 10,10)
    Acos2, Asin2 = potential.scf_compute_coeffs_axi(axi_density1, 10, 10, radial_order=50, costheta_order=50)

    assert numpy.all(numpy.fabs(Acos - Acos2) < 1e-10), \
    "Increasing the radial and costheta order fails for scf_compute_coeffs_axi"

##Tests radial, costheta and phi order from scf_compute_coeffs
def test_density1_Order():
    Acos, Asin = potential.scf_compute_coeffs(density1, 5,5)
    Acos2, Asin2 = potential.scf_compute_coeffs(density1, 5,5, radial_order=19, costheta_order=19, phi_order=19)
    assert numpy.all(numpy.fabs(Acos - Acos2) < 1e-3), \
    "Increasing the radial, costheta, and phi order fails for Acos from scf_compute_coeffs"

    assert numpy.all(numpy.fabs(Asin - Asin) < EPS), \
    "Increasing the radial, costheta, and phi order fails for Asin from scf_compute_coeffs"

## Tests whether scf_compute_axi reduces to scf_compute_spherical for the Hernquist Potential
def test_scf_axiHernquistCoeffs_ReducesToSpherical():
    Aspherical = potential.scf_compute_coeffs_spherical(sphericalHernquistDensity, 10)
    Aaxi = potential.scf_compute_coeffs_axi(sphericalHernquistDensity, 10,10)
    axi_reducesto_spherical(Aspherical,Aaxi, "Hernquist Potential")

## Tests whether scf_compute_axi reduces to scf_compute_spherical for Zeeuw's Potential
def test_scf_axiZeeuwCoeffs_ReducesToSpherical():
    Aspherical = potential.scf_compute_coeffs_spherical(rho_Zeeuw, 10)
    Aaxi = potential.scf_compute_coeffs_axi(rho_Zeeuw, 10,10)
    axi_reducesto_spherical(Aspherical,Aaxi, "Zeeuw Potential")

## Tests whether scf_compute reduces to scf_compute_spherical for Hernquist Potential
def test_scf_HernquistCoeffs_ReducesToSpherical():
    Aspherical = potential.scf_compute_coeffs_spherical(sphericalHernquistDensity, 5)
    Aaxi = potential.scf_compute_coeffs(sphericalHernquistDensity, 5,5)
    reducesto_spherical(Aspherical,Aaxi, "Hernquist Potential")

## Tests whether scf_compute reduces to scf_compute_spherical for Zeeuw's Potential
def test_scf_ZeeuwCoeffs_ReducesToSpherical():
    Aspherical = potential.scf_compute_coeffs_spherical(rho_Zeeuw, 5)
    Aaxi = potential.scf_compute_coeffs(rho_Zeeuw, 5,5,radial_order=20,
                                        costheta_order=20)
    reducesto_spherical(Aspherical,Aaxi, "Zeeuw Potential")

## Tests whether scf density matches with Hernquist density
def test_densMatches_hernquist():
    h = potential.HernquistPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(sphericalHernquistDensity,10)
    scf = SCFPotential()
    assertmsg = "Comparing the density of Hernquist Potential with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(h.dens,scf.dens, assertmsg)

## Tests whether scf density matches with Zeeuw density
def test_densMatches_zeeuw():
    Acos, Asin = potential.scf_compute_coeffs_spherical(rho_Zeeuw,10)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin)
    assertmsg = "Comparing the density of Zeeuw's perfect ellipsoid with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(rho_Zeeuw,scf.dens, assertmsg)

## Tests whether scf density matches with axi_density1
def test_densMatches_axi_density1():
    Acos, Asin = potential.scf_compute_coeffs_axi(axi_density1,50,3)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin)
    assertmsg = "Comparing axi_density1 with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(axi_density1, scf.dens, assertmsg, eps=1e-3)

## Tests whether scf density matches with axi_density2
def test_densMatches_axi_density2():
    Acos, Asin = potential.scf_compute_coeffs_axi(axi_density2,50,3)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin)
    assertmsg = "Comparing axi_density2 with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(axi_density2,scf.dens, assertmsg, eps=1e-3)

## Tests whether scf density matches with NFW
def test_densMatches_nfw():
    nfw = potential.NFWPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(rho_NFW,50, a=50)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin, a=50)
    assertmsg = "Comparing nfw with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(nfw.dens,scf.dens, assertmsg, eps=1e-2)

## Tests whether scf potential matches with Hernquist potential
def test_potentialMatches_hernquist():
    h = potential.HernquistPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(sphericalHernquistDensity,10)
    scf = SCFPotential()
    assertmsg = "Comparing the potential of Hernquist Potential with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(h,scf, assertmsg)

## Tests whether scf Potential matches with NFW
def test_potentialMatches_nfw():
    nfw = potential.NFWPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(rho_NFW,50, a=50)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin, a=50)
    assertmsg = "Comparing nfw with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(nfw,scf, assertmsg, eps=1e-4)

## Tests whether scf Rforce matches with Hernquist Rforce
def test_RforceMatches_hernquist():
    h = potential.HernquistPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(sphericalHernquistDensity,1)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin)
    assertmsg = "Comparing the radial force of Hernquist Potential with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(h.Rforce,scf.Rforce, assertmsg)

## Tests whether scf zforce matches with Hernquist zforce
def test_zforceMatches_hernquist():
    h = potential.HernquistPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(sphericalHernquistDensity,1)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin)
    assertmsg = "Comparing the vertical force of Hernquist Potential with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(h.zforce,scf.zforce, assertmsg)

## Tests whether scf phiforce matches with Hernquist phiforce
def test_phiforceMatches_hernquist():
    h = potential.HernquistPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(sphericalHernquistDensity,1)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin)
    assertmsg = "Comparing the azimuth force of Hernquist Potential with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(h.phiforce,scf.phiforce, assertmsg)

## Tests whether scf Rforce matches with NFW Rforce
def test_RforceMatches_nfw():
    nfw = potential.NFWPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(rho_NFW,50, a=50)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin, a=50)
    assertmsg = "Comparing the radial force of NFW Potential with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(nfw.Rforce,scf.Rforce, assertmsg, eps=1e-3)

## Tests whether scf zforce matches with NFW zforce
def test_zforceMatches_nfw():
    nfw = potential.NFWPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(rho_NFW,50, a=50)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin, a=50)
    assertmsg = "Comparing the vertical force of NFW Potential with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(nfw.zforce,scf.zforce, assertmsg, eps=1e-3)

## Tests whether scf phiforce matches with NFW Rforce
def test_phiforceMatches_nfw():
    nfw = potential.NFWPotential()
    Acos, Asin = potential.scf_compute_coeffs_spherical(rho_NFW,10)
    scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin)
    assertmsg = "Comparing the azimuth force of NFW Potential with SCF fails at R={0}, Z={1}, phi={2}"
    compareFunctions(nfw.phiforce,scf.phiforce, assertmsg)

##############GENERIC FUNCTIONS BELOW###############

##This is used to test whether input as arrays works
def ArrayTest(scf, params):
    def compareFunctions(func, result, i):
        if numpy.isnan(result[i]):
            return numpy.isnan(func(R[i], z[i], phi[i]))
        if numpy.isinf(result[i]):
            return numpy.isinf(func(R[i], z[i], phi[i]))
        return numpy.all(numpy.fabs(result[i] - func(R[i], z[i], phi[i])) < EPS)

    potential = scf(*params).flatten()
    density = scf.dens(*params).flatten()
    Rforce = scf.Rforce(*params).flatten()
    zforce = scf.zforce(*params).flatten()
    phiforce = scf.phiforce(*params).flatten()

    R, z, phi = params
    shape = numpy.array(R*z*phi).shape
    R = (numpy.ones(shape)*R).flatten(); z = (numpy.ones(shape)*z).flatten(); phi = (numpy.ones(shape)*phi).flatten();
    message = "{0} at R={1}, z={2}, phi={3} was found to be {4} where it was expected to be equal to {5}"
    for i in range(len(R)):
        assert compareFunctions(scf, potential, i), \
    message.format("Potential", R[i], z[i], phi[i], potential[i], scf(R[i], z[i], phi[i]))
        assert compareFunctions(scf.dens, density, i), \
    message.format("Density", R[i], z[i], phi[i], density[i], scf.dens(R[i], z[i], phi[i]))
        assert compareFunctions(scf.Rforce, Rforce, i), \
    message.format("Rforce", R[i], z[i], phi[i], Rforce[i], scf.Rforce(R[i], z[i], phi[i]))
        assert compareFunctions(scf.zforce, zforce, i), \
    message.format("zforce", R[i], z[i], phi[i], zforce[i], scf.zforce(R[i], z[i], phi[i]))
        assert compareFunctions(scf.phiforce, phiforce, i), \
    message.format("phiforce", R[i], z[i], phi[i], phiforce[i], scf.phiforce(R[i], z[i], phi[i]))

## This is used to compare scf functions with its corresponding galpy function
def compareFunctions(galpyFunc, scfFunc, assertmsg, Rs=DEFAULT_R, Zs = DEFAULT_Z, phis = DEFAULT_PHI, eps=EPS):
    ##Assert msg must have 3 placeholders ({}) for Rs, Zs, and phis
    for ii in range(len(Rs)):
        for jj in range(len(Zs)):
            for kk in range(len(phis)):
                e = numpy.divide(galpyFunc(Rs[ii],Zs[jj],phis[kk]) - scfFunc(Rs[ii],Zs[jj],phis[kk]), galpyFunc(Rs[ii],Zs[jj],phis[kk]))
                e = numpy.fabs(numpy.fabs(e))
                if galpyFunc(Rs[ii],Zs[jj],phis[kk]) == 0: continue ## Ignoring divide by zero
                assert e < eps, \
                assertmsg.format(Rs[ii],Zs[jj],phis[kk])

##General function that tests whether coefficients for a spherical density has the expected property
def spherical_coeffsTest(Acos, Asin, eps=EPS):
    ## We expect Asin to be zero
    assert Asin is None or numpy.all(numpy.fabs(Asin) <eps), "Confirming Asin = 0 fails"
    ## We expect that the only non-zero values occur at (n,l=0,m=0)
    assert numpy.all(numpy.fabs(Acos[:, 1:, :]) < eps) and numpy.all(numpy.fabs(Acos[:,:,1:]) < eps), \
    "Non Zero value found outside (n,l,m) = (n,0,0)"

##General function that tests whether coefficients for an axi symmetric density has the expected property
def axi_coeffsTest(Acos, Asin):
    ## We expect Asin to be zero
    assert Asin is None or numpy.all(numpy.fabs(Asin) <EPS), "Confirming Asin = 0 fails"
    ## We expect that the only non-zero values occur at (n,l,m=0)
    assert numpy.all(numpy.fabs(Acos[:,:,1:]) < EPS), "Non Zero value found outside (n,l,m) = (n,0,0)"

## Tests whether the coefficients of a spherical density computed using the scf_compute_coeffs_axi reduces to
## The coefficients computed using the scf_compute_coeffs_spherical
def axi_reducesto_spherical(Aspherical,Aaxi,potentialName):
    Acos_s, Asin_s = Aspherical
    Acos_a, Asin_a = Aaxi

    spherical_coeffsTest(Acos_a, Asin_a, eps=1e-10)
    n = min(Acos_s.shape[0], Acos_a.shape[0])
    assert numpy.all(numpy.fabs(Acos_s[:n,0,0] - Acos_a[:n,0,0]) < EPS), \
    "The axi symmetric Acos(n,l=0,m=0) does not reduce to the spherical Acos(n,l=0,m=0) for {0}".format(potentialName)

## Tests whether the coefficients of a spherical density computed using the scf_compute_coeffs reduces to
## The coefficients computed using the scf_compute_coeffs_spherical
def reducesto_spherical(Aspherical,A,potentialName):
    Acos_s, Asin_s = Aspherical
    Acos, Asin = A

    spherical_coeffsTest(Acos, Asin, eps=1e-10)
    n = min(Acos_s.shape[0], Acos.shape[0])
    assert numpy.all(numpy.fabs(Acos_s[:n,0,0] - Acos[:n,0,0]) < EPS), \
    "Acos(n,l=0,m=0) as generated by scf_compute_coeffs does not reduce to the spherical Acos(n,l=0,m=0) for {0}".format(potentialName)

## Hernquist potential as a function of r
def sphericalHernquistDensity(R, z=0, phi=0):
    h = potential.HernquistPotential()
    return h.dens(R,z,phi)

def rho_Zeeuw(R,z,phi,a=1.):
    r, theta, phi = coords.cyl_to_spher(R,z, phi)
    return 3./(4*numpy.pi) * numpy.power((a + r),-4.) * a

def rho_NFW(R, z=0, phi=0.):
    nfw = potential.NFWPotential()
    return nfw.dens(R,z, phi)

def axi_density1(R, z=0, phi=0.):
    r, theta, phi = coords.cyl_to_spher(R,z, phi)
    h = potential.HernquistPotential()
    return h.dens(R, z, phi)*(1 + numpy.cos(theta) + numpy.cos(theta)**2.)

def axi_density2(R, z=0, phi=0.):
    spherical_coords = coords.cyl_to_spher(R,z, phi)
    theta = spherical_coords[1]
    return rho_Zeeuw(R,z,phi)*(1 + numpy.cos(theta) + numpy.cos(theta)**2)

def density1(R, z=0, phi=0.):
    r, theta, phi = coords.cyl_to_spher(R,z, phi)
    h = potential.HernquistPotential(2)
    return h.dens(R, z, phi)*(1 + numpy.cos(theta) + numpy.cos(theta)**2.)*(1 + numpy.cos(phi) + numpy.sin(phi))

max_stars_repo_name: morganb-phys/galpy

max_issues_repo_name: morganb-phys/galpy

max_forks_repo_name: morganb-phys/galpy


text:
import sqlite3
import os
import time
import json
import configparser
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import font_manager, colors
from tqdm import trange, tqdm
from collections import Counter
from wordcloud import WordCloud

# 载入数据库名
CONFIG = configparser.ConfigParser()
CONFIG.read('config.ini', encoding='utf-8')
DB_NAME = CONFIG['common'].get('database_name', fallback='data')

# 连接数据库
conn = sqlite3.connect('data/' + DB_NAME + '.sqlite3')
conn.row_factory = sqlite3.Row
cursor = conn.cursor()

if not os.path.exists('analysize'):
    os.mkdir('analysize')

def str_to_timestamp(time_str):
    '''将时间字符串转换成时间戳'''
    return int(time.mktime(time.strptime(time_str, "%Y-%m-%d %H:%M:%S")))

def pic01():
    '''统计视频信息按时段的分布'''

    print('pic0 and pic1')

    # 数据库查询
    view_sum = [0] * 24
    count = [0] * 24
    for day in trange(1, 32, position=1):
        for hour in trange(0, 24, position=0):
            st = str_to_timestamp('2021-08-{} {}:00:00'.format(day, hour))
            ed = str_to_timestamp('2021-08-{} {}:59:59'.format(day, hour))
            count[hour] += cursor.execute(
                'SELECT COUNT() from VIDEO WHERE pubdate BETWEEN ? and ?', (st, ed)).fetchone()[0]
            view_sum[hour] += cursor.execute(
                'SELECT SUM(view) from VIDEO WHERE pubdate BETWEEN ? and ?', (st, ed)).fetchone()[0]

    # 数据准备
    N = 24
    view_sum = np.array(view_sum)
    count = np.array(count)
    theta = np.linspace(0.0 + np.pi / N, 2 * np.pi +
                        np.pi / N, N, endpoint=False)
    width = 2 * np.pi / N
    colors = np.array(['#7bbfea', '#33a3dc'] * 12)

    # 图表绘制
    plt.figure(figsize=(16, 9))
    plt.style.use('ggplot')
    ax = plt.subplot(111, projection='polar')
    ax.set_axisbelow(True)
    ax.set_theta_direction(-1)
    ax.set_theta_zero_location('N')
    ax.set_rlabel_position(90)
    ax.bar(theta, count / 1000, width=width, bottom=0.0, color=colors)
    ax.set_xticks(np.linspace(0, 2*np.pi, 24, endpoint=False))
    ax.set_xticklabels(range(0, 24), fontsize=20)
    ax.set_yticks([i for i in range(0, 6)])
    ax.set_yticklabels([0] + ['{}K'.format(i)
                       for i in range(1, 6)], fontsize=15)
    font = font_manager.FontProperties(fname='font/msyh.ttc', size=28)
    ax.set_title('2021年08月 B站数码区视频投稿数量在一天中的分布', fontproperties=font)
    plt.suptitle("2021.09.05 统计", x=0.75, y=0.9,
                 fontproperties=font, fontsize=18)
    plt.savefig('analysize/pic0.png')

    plt.figure(figsize=(16, 9))
    plt.style.use('ggplot')
    ax = plt.subplot(111, projection='polar')
    ax.set_axisbelow(True)
    ax.set_theta_direction(-1)
    ax.set_theta_zero_location('N')
    ax.set_rlabel_position(90)
    ax.bar(theta, view_sum / count / 1000,
           width=width, bottom=0.0, color=colors)
    ax.set_xticks(np.linspace(0, 2*np.pi, 24, endpoint=False))
    ax.set_xticklabels(range(0, 24), fontsize=20)
    ax.set_yticks([i for i in range(0, 11, 2)])
    ax.set_yticklabels([0] + ['{}K'.format(i)
                       for i in range(2, 11, 2)], fontsize=15)
    font = font_manager.FontProperties(fname='font/msyh.ttc', size=28)
    plt.title('2021年08月 B站数码区一天中各时段发布视频的平均播放量', fontproperties=font)
    plt.suptitle("2021.09.05 统计", x=0.75, y=0.9,
                 fontproperties=font, fontsize=18)
    plt.savefig('analysize/pic1.png')

def pic2():
    '''分析标签'''

    print('\npic2')

    # 数据库查询与数据统计
    st = str_to_timestamp('2021-08-{} {}:00:00'.format(1, 0))
    ed = str_to_timestamp('2021-08-{} {}:59:59'.format(31, 23))
    total_num = cursor.execute(
        'SELECT COUNT() from VIDEO WHERE pubdate BETWEEN ? and ?', (st, ed)).fetchone()[0]
    data = cursor.execute(
        'SELECT * from VIDEO WHERE pubdate BETWEEN ? and ?', (st, ed))

    count = Counter()
    with tqdm(total=total_num) as pbar:
        for row in data:
            keywords = row['keywords']
            count += Counter(keywords.split(','))
            pbar.update(1)
    with open('count.json', 'w') as f:
        json.dump(count, f)

    # 删除无效标签
    del_list = ['数码', '科技', '哔哩哔哩', 'B站', '弹幕',
                'Bilibili', '打卡挑战', '科技猎手', '必剪创作']
    for name in del_list:
        del count[name]

    N = 30

    # 生成词云
    color_list = ['#2585a6', '#228fbd', '#2468a2',
                  '#2570a1', '#009ad6', '#145b7d']
    colormap = colors.ListedColormap(color_list)
    wc = WordCloud(width=1600, height=900, font_path='font/msyh.ttc', max_words=N, mode='RGBA',
                   background_color='white', colormap=colormap).generate_from_frequencies(dict(count.most_common(N)))
    wc.to_file('analysize/pic2-1.png')

    # 生成图表
    y, x = map(list, zip(*count.most_common(N)))

    plt.figure(figsize=(16, 9))
    plt.style.use('ggplot')
    font = font_manager.FontProperties(fname='font/msyh.ttc')
    y_pos = range(len(y), 0, -1)
    plt.barh(y_pos, x, color='#33a3dc')
    plt.yticks(y_pos, y, fontproperties=font, fontsize=14, rotation=10)
    plt.xticks(fontsize=18)
    plt.xlabel('词频 / 次', fontproperties=font, fontsize=20)
    plt.ylabel('标签', fontproperties=font, fontsize=20)
    plt.title('2021年08月 B站数码区视频标签词频统计前 30 名\n',
              fontproperties=font, fontsize=28)
    plt.suptitle("2021.09.05 统计", x=0.75, y=0.88,
                 fontproperties=font, fontsize=18)
    plt.tight_layout()
    plt.savefig('analysize/pic2-0.png')

def pic3():
    '''分析播放量与视频时长关系'''

    print('\npic3')

    # 查询数据库
    x = []
    y = []
    interval = 600
    N = 6
    for i in trange(N):
        y.append(cursor.execute('SELECT AVG(view) from VIDEO WHERE duration BETWEEN ? and ?',
                 (interval * i, interval * (i + 1) - 1)).fetchone()[0])
        x.append('{}~{} 分钟'.format(10 * i, 10 * (i + 1)))
    x.append('{} 分钟以上'.format(10 * N))
    y.append(cursor.execute(
        'SELECT AVG(view) from VIDEO WHERE duration > ?', (interval * N,)).fetchone()[0])

    # 生成图表
    plt.figure(figsize=(16, 9))
    plt.style.use('ggplot')
    font = font_manager.FontProperties(fname='font/msyh.ttc')
    plt.bar(range(len(x)), y, color='#33a3dc')
    plt.xticks(range(len(x)), x, fontproperties=font, fontsize=18)
    plt.yticks(fontsize=18)
    plt.xlabel('视频时长', fontproperties=font, fontsize=20)
    plt.ylabel('平均播放量 / 次', fontproperties=font, fontsize=20)
    plt.title('2021年08月 B站数码区视频平均播放量随视频总时长的分布\n',
              fontproperties=font, fontsize=28)
    plt.suptitle("2021.09.05 统计", x=0.75, y=0.92,
                 fontproperties=font, fontsize=18)
    plt.savefig('analysize/pic3.png')

pic01()
pic2()
pic3()


text:
"""Stereo3D module"""
import os
import sys
import time
import traceback
import numpy as np
import cv2
from pymsgbox import prompt, alert
from pyntcloud import PyntCloud
try:
    from i3drsgm import I3DRSGM
    I3DRSGM_IMPORTED = True
except ImportError:
    I3DRSGM = None
    I3DRSGM_IMPORTED = False
import pandas as pd

class Stereo3D():
    def __init__(self, stereo_camera, stereo_calibration,
                 stereo_matcher=None, show_window=True):
        """
        Initialisation function for PickPlace3D. Defines the devices used.
        If you would like to use the methods without connecting
        a camera then initaialise Stereo3D as 's3D = Stereo3D()'.
        :param left_cal_file:
            filepath to left image calibration file (e.g. left.yaml)
        :param right_cal_file:
            filepath to right image calibration file (e.g. right.yaml)
        :param stereo_camera:
            stereo camera used for generating 3D and 2D detection
        :type left_cal_file: string
        :type right_cal_file: string
        :type stereo_camera: StereoCapture.StereoCapture
        """
        self.change_camera(stereo_camera)
        self.stereo_calibration = stereo_calibration
        self.Q = self.stereo_calibration.stereo_cal["q"]

        self.image_left = None
        self.image_right = None
        self.rect_image_left = None
        self.rect_image_right = None
        self.disparity = None
        self.depth = None

        self.EXIT_CODE_QUIT = -1
        self.EXIT_CODE_GRAB_3D_SUCCESS = 1
        self.EXIT_CODE_FAILED_TO_GRAB_3D = -2

        self.save_index = 0

        self.cv_window_name_Controls = "[Stereo3D] Controls"
        self.cv_window_name_Images = "[Stereo3D] Images"

        self.show_window = show_window
        if self.show_window:
            cv2.namedWindow(self.cv_window_name_Controls, cv2.WINDOW_NORMAL)
            # cv2.namedWindow(self.cv_window_name_Images, cv2.WINDOW_NORMAL)

            cv2.setMouseCallback(
                self.cv_window_name_Images, self.on_window_mouse)

            cv2.resizeWindow(self.cv_window_name_Controls, 400, 0)

        self.matcher_name = ""
        self.change_matcher(stereo_matcher)

        # Once init has been called matcher can be changed e.g.
        # s3D = Stereo3D()
        # s3D.ply_header = "ply..."
        self.ply_header = (
            "ply\n"
            "format ascii 1.0\n"
            "element vertex %(vert_num)d\n"
            "property float x\n"
            "property float y\n"
            "property float z\n"
            "property uchar red\n"
            "property uchar green\n"
            "property uchar blue\n"
            "end_header\n")

        self.frame_count = 0

    def change_camera(self, stereo_camera):
        self.stereo_camera = stereo_camera

    def change_matcher(self, stereo_matcher):
        default_min_disp = 1000
        default_num_disparities = 20
        default_block_size = 12
        default_uniqueness_ratio = 15
        default_texture_threshold = 15
        default_speckle_size = 0
        default_speckle_range = 500

        if (stereo_matcher is None):
            self.matcher = cv2.StereoBM_create()
            calc_block = (2 * default_block_size + 5)
            self.matcher.setBlockSize(calc_block)
            self.matcher.setMinDisparity(int(default_min_disp - 1000))
            self.matcher.setNumDisparities(16*(default_num_disparities+1))
            self.matcher.setUniquenessRatio(default_uniqueness_ratio)
            self.matcher.setTextureThreshold(default_texture_threshold)
            self.matcher.setSpeckleWindowSize(default_speckle_size)
            self.matcher.setSpeckleRange(default_speckle_range)
        else:
            if stereo_matcher == "BM":
                self.matcher_name = stereo_matcher
                self.matcher = cv2.StereoBM_create()
                calc_block = (2 * default_block_size + 5)
                self.matcher.setBlockSize(calc_block)
                self.matcher.setMinDisparity(int(default_min_disp - 1000))
                self.matcher.setNumDisparities(16*(default_num_disparities+1))
                self.matcher.setUniquenessRatio(default_uniqueness_ratio)
                self.matcher.setTextureThreshold(default_texture_threshold)
                self.matcher.setSpeckleWindowSize(default_speckle_size)
                self.matcher.setSpeckleRange(default_speckle_range)
            elif stereo_matcher == "SGBM":
                self.matcher_name = stereo_matcher
                self.matcher = cv2.StereoSGBM_create()
                calc_block = (2 * default_block_size + 5)
                self.matcher.setBlockSize(calc_block)
                self.matcher.setMinDisparity(int(default_min_disp - 1000))
                self.matcher.setNumDisparities(16*(default_num_disparities+1))
                self.matcher.setUniquenessRatio(default_uniqueness_ratio)
                # self.matcher.setTextureThreshold(default_texture_threshold)
                self.matcher.setSpeckleWindowSize(default_speckle_size)
                self.matcher.setSpeckleRange(default_speckle_range)
            elif stereo_matcher == "I3DRSGM":
                if I3DRSGM_IMPORTED:
                    self.matcher_name = stereo_matcher
                    self.matcher = I3DRSGM()
                    self.matcher.setWindowSize(11)
                    self.matcher.setMinDisparity(0)
                    self.matcher.setDisparityRange(16*120)
                    self.matcher.setPyamidLevel(6)
                else:
                    raise Exception("Failed to import I3DRSGM")
            else:
                self.matcher = stereo_matcher

        if self.show_window:
            cv2.destroyWindow(self.cv_window_name_Controls)
            cv2.namedWindow(self.cv_window_name_Controls, cv2.WINDOW_NORMAL)

            cv2.createTrackbar(
                "Min disp", self.cv_window_name_Controls,
                default_min_disp, 2000, self.on_min_disparity_trackbar)
            cv2.createTrackbar(
                "Disp", self.cv_window_name_Controls,
                default_num_disparities, 30, self.on_num_disparities_trackbar)
            cv2.createTrackbar(
                "Blck sze", self.cv_window_name_Controls,
                default_block_size, 100, self.on_block_size_trackbar)

            cv2.createTrackbar(
                "Uniq", self.cv_window_name_Controls,
                default_uniqueness_ratio, 100,
                self.on_uniqueness_ratio_trackbar)
            if stereo_matcher == "BM":
                cv2.createTrackbar(
                    "Texture", self.cv_window_name_Controls,
                    default_texture_threshold, 100,
                    self.on_texture_threshold_trackbar)

            cv2.createTrackbar(
                "Sp size", self.cv_window_name_Controls,
                default_speckle_size, 500, self.on_speckle_size_trackbar)
            cv2.createTrackbar(
                "Sp range", self.cv_window_name_Controls,
                default_speckle_range, 1000, self.on_speckle_range_trackbar)

    def connect(self):
        """
        Connect to devices needed for the pick and place process.
        :returns: success
        :rtype: bool
        """
        res = self.stereo_camera.connect()
        return res

    def gen3D(self, left_image, right_image):
        if self.matcher_name == "I3DRSGM":
            if I3DRSGM_IMPORTED:
                _, disparity = self.matcher.forwardMatch(
                    left_image, right_image
                )
                disparity = -disparity.astype(np.float32)
                disparity[disparity == 99999] = 0.0
                disparity[disparity == 0] = int(disparity.min() - 1)
                disparity = np.nan_to_num(
                    disparity, nan=0.0, posinf=0.0, neginf=0.0)
            else:
                raise Exception("Failed to import I3DRSGM")
        else:
            disparity = self.matcher.compute(left_image, right_image)
            disparity = disparity.astype(np.float32) / 16.0
        return disparity

    def genDepth(self, disparity):
        depth = cv2.reprojectImageTo3D(disparity, self.Q)
        return depth

    def write_ply(self, filename, disp, depth, image):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = np.float32(image)
        image = image/255

        mask = disp > disp.min()
        points = depth[mask]
        image = image[mask]

        # TODO replace this with better filter in-case
        # 3D is every large an this becomes invalid
        points[points >= 100] = 0
        points[points <= -100] = 0

        points = points.reshape(-1, 3)
        image = image.reshape(-1, 3)

        points = np.hstack((points, image))

        cloud = PyntCloud(pd.DataFrame(
            # same arguments that you are passing to visualize_pcl
            data=points,
            columns=["x", "y", "z", "red", "green", "blue"]))

        cloud.to_file(filename)

    def scale_disparity(self, disparity):
        minV, maxV, _, _ = cv2.minMaxLoc(disparity)
        if maxV - minV != 0:
            scaled_disp = cv2.convertScaleAbs(
                disparity, alpha=255.0/(maxV - minV),
                beta=-minV * 255.0/(maxV - minV))
            return scaled_disp
        else:
            return np.zeros(disparity.shape, np.uint8)

    def on_min_disparity_trackbar(self, val):
        min_disp = int(val - 1000)
        self.matcher.setMinDisparity(min_disp)

    def on_block_size_trackbar(self, val):
        if self.matcher_name == "I3DRSGM":
            if I3DRSGM_IMPORTED:
                self.matcher.setWindowSize(2 * val + 5)
            else:
                raise Exception("Failed to import I3DRSGM")
        else:
            self.matcher.setBlockSize(2 * val + 5)

    def on_num_disparities_trackbar(self, val):
        if self.matcher_name == "I3DRSGM":
            if I3DRSGM_IMPORTED:
                self.matcher.setDisparityRange(16*((val*10)+1))
            else:
                raise Exception("Failed to import I3DRSGM")
        else:
            self.matcher.setNumDisparities(16*(val+1))

    def on_texture_threshold_trackbar(self, val):
        if self.matcher_name != "I3DRSGM":
            self.matcher.setTextureThreshold(val)

    def on_uniqueness_ratio_trackbar(self, val):
        if self.matcher_name != "I3DRSGM":
            self.matcher.setUniquenessRatio(val)

    def on_speckle_size_trackbar(self, val):
        if self.matcher_name != "I3DRSGM":
            self.matcher.setSpeckleWindowSize(val)

    def on_speckle_range_trackbar(self, val):
        if (self.matcher_name != "I3DRSGM"):
            self.matcher.setSpeckleRange(val)

    def on_window_mouse(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            print("click")

    def grab3D(self, isRectified=False):
        res, image_left, image_right = self.stereo_camera.grab()
        if res:
            self.image_left = image_left
            self.image_right = image_right
            if isRectified:
                self.rect_image_left = self.image_left
                self.rect_image_right = self.image_right
            else:
                self.rect_image_left, self.rect_image_right = \
                    self.stereo_calibration.rectify_pair(
                        image_left, image_right)

            disp = self.gen3D(self.rect_image_left, self.rect_image_right)
            self.disparity = disp

            return True, disp
        else:
            print("Failed to grab stereo image pair")
            return False, None

    def save_point_cloud(self, disparity, image,
                         defaultSaveFolder="", points_file_string="output.ply",
                         confirm_folder=True):
        # prompt user for save location
        resp = defaultSaveFolder
        if confirm_folder:
            try:
                resp = prompt(
                    text='Saving 3D Point Cloud to path: ',
                    title='Save 3D Point Cloud', default=defaultSaveFolder)
            except AssertionError:
                _, _, tb = sys.exc_info()
                traceback.print_tb(tb)  # Fixed format
                tb_info = traceback.extract_tb(tb)
                filename, line, func, text = tb_info[-1]

                err_msg = 'An error occurred on line {} in statement {}'
                err_msg = err_msg.format(line, text)
                print(err_msg)
                err_msg = 'Will use default save folder: {}'
                err_msg = err_msg.format(defaultSaveFolder)
                print(err_msg)
        if (resp is not None and disparity is not None and image is not None):
            # define name of output point cloud ply file
            ply_filename = os.path.join(resp, points_file_string)

            # generate depth from disparity
            print("Generating depth from disparity...")
            depth = self.genDepth(disparity)
            print("Saving point cloud...")
            # write 3D data to ply with color from image on points
            self.write_ply(ply_filename, disparity, depth, image)
            print("Point cloud save complete.")
            if confirm_folder:
                alert('3D point cloud saved.', 'Save 3D Point Cloud')
        else:
            print("invalid prompt response or disparity/image is empty")

    def save_all_current(self, saveFolder):
        confirm_folder = False
        left_file_string = str(self.save_index)+"_l.png"
        right_file_string = str(self.save_index)+"_r.png"
        self.stereo_camera.save_images(
            self.image_left, self.image_right,
            saveFolder, left_file_string, right_file_string,
            confirm_folder)

        left_file_string = "rect_"+str(self.save_index)+"_l.png"
        right_file_string = "rect_"+str(self.save_index)+"_r.png"
        self.stereo_camera.save_images(
            self.rect_image_left, self.rect_image_right,
            saveFolder, left_file_string, right_file_string,
            confirm_folder)

        points_file_string = "points_"+str(self.save_index)+".ply"
        self.save_point_cloud(
            self.disparity, self.rect_image_left,
            saveFolder, points_file_string, confirm_folder)

    def run_frame_no_gui(self, saveFolder, isRectified=False):
        # grab 3D disparity from stereo camera
        res, _ = self.grab3D(isRectified)
        if res:
            self.save_all_current(saveFolder)
        return res

    def run_frame(self, defaultSaveFolder="", isRectified=False,
                  confirm_folder=True, colormap=cv2.COLORMAP_JET):
        # grab 3D disparity from stereo camera
        res, disp = self.grab3D(isRectified)
        if res:
            self.frame_count += 1
            # prepare images for displaying
            display_image = np.zeros((640, 480), np.uint8)

            rect_image_left_resized = \
                self.stereo_camera.image_resize(
                    self.rect_image_left, height=640)
            rect_image_right_resized = \
                self.stereo_camera.image_resize(
                    self.rect_image_right, height=640)

            disp_resized = self.scale_disparity(
                self.stereo_camera.image_resize(disp, height=640))
            disp_black_mask = disp_resized <= 0
            # apply color map to disparity
            disp_colormap = cv2.applyColorMap(disp_resized, colormap)
            disp_colormap[disp_black_mask != 0] = [0, 0, 0]

            left_right_dual_gray = np.concatenate(
                (rect_image_left_resized, rect_image_right_resized), axis=1)
            left_right_dual = cv2.cvtColor(
                left_right_dual_gray, cv2.COLOR_GRAY2RGB)

            (lr_dual_h, lr_dual_w, _) = left_right_dual.shape
            (d_h, d_w, _) = disp_colormap.shape

            spacer_width_raw = (lr_dual_w - d_w)
            if (spacer_width_raw % 2) == 0:
                # even
                spacer_width_1 = int((spacer_width_raw / 2))
                spacer_width_2 = int((spacer_width_raw / 2))
            else:
                # odd
                spacer_width_1 = int((spacer_width_raw / 2))
                spacer_width_2 = int((spacer_width_raw / 2) + 1)

            disp_spacer_1 = np.zeros((640, spacer_width_1, 3), np.uint8)
            disp_spacer_2 = np.zeros((640, spacer_width_2, 3), np.uint8)
            disp_spaced = np.concatenate(
                (disp_spacer_1, disp_colormap, disp_spacer_2), axis=1)

            display_image = np.concatenate(
                (disp_spaced, left_right_dual), axis=0)
            display_image_resize = self.stereo_camera.image_resize(
                display_image, height=640)

            font = cv2.FONT_HERSHEY_DUPLEX
            bottomLeftCornerOfText = (10, 20)
            fontScale = 0.4
            fontColor = (255, 255, 255)
            lineType = 1

            cv2.putText(
                display_image_resize, 'Frame: {}'.format(self.frame_count),
                bottomLeftCornerOfText,
                font,
                fontScale,
                fontColor,
                lineType)

            if self.show_window:
                # display disparity with stereo images
                cv2.imshow(self.cv_window_name_Images, display_image_resize)

        k = cv2.waitKey(1)
        if k == ord('q'):  # exit if 'q' key pressed
            return self.EXIT_CODE_QUIT
        elif k == ord('s'):  # save stereo image pair
            left_file_string = str(self.save_index)+"_l.png"
            right_file_string = str(self.save_index)+"_r.png"
            self.stereo_camera.save_images(
                self.image_left, self.image_right,
                defaultSaveFolder, left_file_string, right_file_string,
                confirm_folder)
            self.save_index += 1
        elif k == ord('r'):  # save rectified stereo image pair
            left_file_string = "rect_"+str(self.save_index)+"_l.png"
            right_file_string = "rect_"+str(self.save_index)+"_r.png"
            self.stereo_camera.save_images(
                self.rect_image_left, self.rect_image_right,
                defaultSaveFolder, left_file_string, right_file_string,
                confirm_folder)
            self.save_index += 1
        elif k == ord('p'):  # save 3D data as point cloud
            points_file_string = "points_"+str(self.save_index)+".ply"
            self.save_point_cloud(
                self.disparity, self.rect_image_left,
                defaultSaveFolder, points_file_string,
                confirm_folder)
            self.save_index += 1
        elif k == ord('1'):  # change tp OpenCV BM
            self.change_matcher("BM")
        elif k == ord('2'):  # change to OpenCV SGBM
            self.change_matcher("SGBM")
        if self.show_window:
            window_image_prop = cv2.getWindowProperty(
                self.cv_window_name_Images, cv2.WND_PROP_VISIBLE)
            window_control_prop = cv2.getWindowProperty(
                self.cv_window_name_Images, cv2.WND_PROP_VISIBLE)
            if window_image_prop < 1:
                return self.EXIT_CODE_QUIT
            if window_control_prop < 1:
                return self.EXIT_CODE_QUIT

        if res:
            return self.EXIT_CODE_GRAB_3D_SUCCESS
        else:
            return self.EXIT_CODE_FAILED_TO_GRAB_3D

    def run(self, defaultSaveFolder="", isRectified=False,
            frame_delay=0, confirm_folder=True, colormap=cv2.COLORMAP_JET):
        # connect to stereo camera
        connected = False
        while(not connected):
            connected = self.connect()
            time.sleep(1)
        if self.show_window:
            cv2.namedWindow(self.cv_window_name_Images, cv2.WINDOW_NORMAL)
        while(True):
            exit_code = self.run_frame(
                defaultSaveFolder, isRectified, confirm_folder, colormap)
            if (exit_code == self.EXIT_CODE_QUIT):
                break
            if (exit_code == self.EXIT_CODE_FAILED_TO_GRAB_3D):
                time.sleep(1)
            time.sleep(frame_delay)

        self.stereo_camera.close()