text:
from shfl.model.model import TrainableModel
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn import metrics

class LinearClassifierModel(TrainableModel):
    """
    This class offers support for scikit-learn linear classification models. By default, LogisticRegression is used.
    It implements [TrainableModel](../Model/#trainablemodel-class)

    # Arguments:
        n_features: integer number of features (independent variables).
        classes: array of classes to predict. At least 2 classes must be provided.
        model: Optional. Sklearn Linear Model instance to use. If it is not provided, a LogisticRegression instance
            will be used. It has been tested with LogisticRegression and LinearSVC instances but it should work for
            every linear model defined by intercept_ and coef_ attributes.
    """
    def __init__(self, n_features, classes, model=None):
        if model is None:
            model = LogisticRegression(solver='lbfgs', multi_class='auto')
        self._check_initialization(n_features, classes)
        self._model = model
        self._n_features = n_features
        classes = np.sort(np.asarray(classes))
        self._model.classes_ = classes
        n_classes = len(classes)
        if n_classes == 2:
            n_classes = 1
        self.set_model_params([np.zeros(n_classes), np.zeros((n_classes, n_features))])

    def train(self, data, labels):
        """
        Implementation of abstract method of class [TrainableModel](../Model/#trainablemodel-class)

        # Arguments
            data: Data, array-like of shape (n_samples, n_features)
            labels: Target classes, array-like of shape (n_samples,)
        """
        self._check_data(data)
        self._check_labels_train(labels)

        self._model.fit(data, labels)

    def predict(self, data):
        """
        Implementation of abstract method of class [TrainableModel](../Model/#trainablemodel-class)

        Arguments:
            data: Data, array-like of shape (n_samples, n_features)
        """

        return self._model.predict(data)

    def evaluate(self, data, labels):
        """
        Implementation of abstract method of class [TrainableModel](../Model/#trainablemodel-class)
        Metrics for evaluating model's performance.

        Arguments:
            data: Data, array-like of shape (n_samples, n_features)
            labels: Target classes, array-like of shape (n_samples,)
        """
        self._check_data(data)
        self._check_labels_predict(labels)

        prediction = self.predict(data)
        bas = metrics.balanced_accuracy_score(labels, prediction)
        cks = metrics.cohen_kappa_score(labels, prediction)

        return bas, cks

    def performance(self, data, labels):
        """
        Implementation of abstract method of class [TrainableModel](../Model/#trainablemodel-class)

        Arguments:
            data: Data, array-like of shape (n_samples, n_features)
            labels: Target classes, array-like of shape (n_samples,)
        """
        self._check_data(data)
        self._check_labels_predict(labels)

        prediction = self.predict(data)
        bas = metrics.balanced_accuracy_score(labels, prediction)

        return bas

    def get_model_params(self):
        """
        Implementation of abstract method of class [TrainableModel](../Model/#trainablemodel-class)
        """

        return [self._model.intercept_, self._model.coef_]

    def set_model_params(self, params):
        """
        Implementation of abstract method of class [TrainableModel](../Model/#trainablemodel-class)
        """

        self._model.intercept_ = params[0]
        self._model.coef_ = params[1]

    def _check_data(self, data):
        """
        Method that checks whether the data dimension is correct.
        """
        if data.ndim == 1:
            if self._n_features != 1:
                raise AssertionError("Data need to have the same number of features described by the model, " + str(self._n_features)
                                     + ". Current data have only 1 feature.")
        elif data.shape[1] != self._n_features:
            raise AssertionError("Data need to have the same number of features described by the model, " + str(self._n_features) +
                                 ". Current data has " + str(data.shape[1]) + " features.")

    def _check_labels_train(self, labels):
        """
        Method that checks whether the classes are correct.
        When training, the classes in client's data must be the same as the input ones.

        # Arguments:
            labels: array with classes
        """
        classes = np.unique(np.asarray(labels))
        if not np.array_equal(self._model.classes_, classes):
            raise AssertionError("When training, labels need to have the same classes described by the model, "
                                 + str(self._model.classes_) + ". Labels of this node are " + str(classes) + " .")

    def _check_labels_predict(self, labels):
        """
        Method that checks whether the classes are correct.
        When predicting, the classes in data must be a subset of the trained ones.

        # Arguments:
            labels: array with classes
        """
        classes = np.unique(np.asarray(labels))
        if not set(classes) <= set(self._model.classes_):
            raise AssertionError("When predicting, labels need to be a subset of the classes described by the model, " + str(self._model.classes_)
                                 + ". Labels in the given data are " + str(classes) + " .")

    @staticmethod
    def _check_initialization(n_features, classes):
        """
        Method that checks if model's initialization is correct.
        The number of features must be an integer equal or greater to one, and there must be at least two classes.

        # Arguments:
            n_features: number of features
            classes: array of classes to predict
        """
        if not isinstance(n_features, int):
            raise AssertionError("n_features must be a positive integer number. Provided " + str(n_features) + " features.")
        if n_features < 0:
            raise AssertionError("It must verify that n_features > 0. Provided value " + str(n_features) + ".")
        if len(classes) < 2:
            raise AssertionError("It must verify that the number of classes > 1. Provided " + str(len(classes)) + " classes.")
        if len(np.unique(classes)) != len(classes):
            classes = list(classes)
            duplicated_classes = [i_class for i_class in classes if classes.count(i_class) > 1]
            raise AssertionError("No duplicated classes allowed. Class(es) duplicated: " + str(duplicated_classes) )


text:
import numpy as _np
from ..abstract_sampler import AbstractSampler

from netket.legacy.stats import mean as _mean

from numba import jit
from numba import int64, float64
from netket.legacy import random as _random
import math

class MetropolisHastingsPt(AbstractSampler):
    def __init__(
        self,
        machine,
        kernel,
        n_replicas=32,
        sweep_size=None,
    ):
        super().__init__(machine)

        self.n_replicas = n_replicas

        self.sweep_size = sweep_size

        self._kernel = kernel

        self.machine_pow = 2.0

        # Contains quantities to compute diffusion coefficient of replicas
        # beta_stats[0] = position of replica at beta=1 (do not reset this!)
        # beta_stats[1] = running average of beta=1 position
        # beta_stats[2] = running average of beta=1 position**2
        # beta_stats[3] = current number of exchange steps performed
        self._beta_stats = _np.zeros(4)

        self.reset(True)

    @property
    def n_replicas(self):
        return self._n_replicas

    @n_replicas.setter
    def n_replicas(self, n_replicas):
        if n_replicas < 1:
            raise ValueError("Expected n_replicas>0. ")

        self._n_replicas = n_replicas

        self._state = _np.zeros((n_replicas, self._input_size))
        self._state1 = _np.copy(self._state)

        self._log_values = _np.zeros(n_replicas, dtype=_np.complex128)
        self._log_values_1 = _np.zeros(n_replicas, dtype=_np.complex128)
        self._log_prob_corr = _np.zeros(n_replicas)

        # Linearly spaced inverse temperature
        self._beta = _np.empty(n_replicas)

        for i in range(n_replicas):
            self._beta[i] = 1.0 - float(i) / float(n_replicas)

        # some temporary arrays
        self._proposed_beta = _np.empty(n_replicas)
        self._beta_prob = _np.empty(n_replicas)

    @property
    def machine_pow(self):
        return self._machine_pow

    @machine_pow.setter
    def machine_pow(self, m_power):
        self._machine_pow = m_power

    @property
    def sweep_size(self):
        return self._sweep_size

    @sweep_size.setter
    def sweep_size(self, sweep_size):
        self._sweep_size = sweep_size if sweep_size != None else self._input_size
        if self._sweep_size < 0:
            raise ValueError("Expected a positive integer for sweep_size ")

    def reset(self, init_random=False):
        if init_random:
            self._kernel.random_state(self._state)

        self._log_values = self.machine.log_val(self._state, out=self._log_values)

        self._accepted_samples = _np.zeros(self._n_replicas)
        self._total_samples = 0

        self._beta_stats[1:].fill(0)

    @staticmethod
    @jit(nopython=True)
    def _fixed_beta_acceptance_kernel(
        state,
        state1,
        log_values,
        log_values_1,
        log_prob_corr,
        machine_pow,
        beta,
        accepted,
    ):

        for i in range(state.shape[0]):
            prob = math.exp(
                machine_pow
                * beta[i]
                * (log_values_1[i] - log_values[i] + log_prob_corr[i]).real
            )

            assert not math.isnan(prob)

            accept = prob > _random.uniform(0, 1)

            if accept:
                log_values[i] = log_values_1[i]
                state[i] = state1[i]
                accepted[i] += 1

    def __next__(self):

        _log_val = self.machine.log_val
        _fb_acc_kernel = self._fixed_beta_acceptance_kernel
        _state = self._state
        _state1 = self._state1
        _log_values = self._log_values
        _log_values_1 = self._log_values_1
        _log_prob_corr = self._log_prob_corr
        _machine_pow = self._machine_pow
        _accepted_samples = self._accepted_samples
        _t_kernel = self._kernel.transition
        _beta = self._beta
        _beta_stats = self._beta_stats
        _proposed_beta = self._proposed_beta
        _beta_prob = self._beta_prob

        for sweep in range(self.sweep_size):

            # Propose a new state using the transition kernel
            _t_kernel(_state, _state1, _log_prob_corr)

            _log_values_1 = _log_val(_state1, out=_log_values_1)

            # Acceptance Kernel for fixed-beta moves
            _fb_acc_kernel(
                _state,
                _state1,
                _log_values,
                _log_values_1,
                _log_prob_corr,
                _machine_pow,
                _beta,
                _accepted_samples,
            )

            # Transition + Acceptance Kernel for replica exchange moves
            self._exchange_step_kernel(
                _log_values,
                _machine_pow,
                _beta,
                _proposed_beta,
                _beta_prob,
                _beta_stats,
                _accepted_samples,
            )

        self._total_samples += self.sweep_size
        return self._state[_np.intp(_beta_stats[0])].reshape(1, -1)

    @staticmethod
    @jit(nopython=True)
    def _exchange_step_kernel(
        log_values, machine_pow, beta, proposed_beta, prob, beta_stats, accepted_samples
    ):
        # Choose a random swap order (odd/even swap)
        swap_order = _random.randint(0, 2, size=()).item()

        n_replicas = beta.shape[0]

        for i in range(swap_order, n_replicas, 2):
            inn = (i + 1) % n_replicas
            proposed_beta[i] = beta[inn]
            proposed_beta[inn] = beta[i]

        for i in range(n_replicas):
            prob[i] = math.exp(
                machine_pow * (proposed_beta[i] - beta[i]) * log_values[i].real
            )

        for i in range(swap_order, n_replicas, 2):
            inn = (i + 1) % n_replicas

            prob[i] *= prob[inn]

            if prob[i] > _random.uniform(0, 1):
                # swapping status
                beta[i], beta[inn] = beta[inn], beta[i]
                accepted_samples[i], accepted_samples[inn] = (
                    accepted_samples[inn],
                    accepted_samples[i],
                )

                if beta_stats[0] == i:
                    beta_stats[0] = inn
                elif beta_stats[0] == inn:
                    beta_stats[0] = i

        # Update statistics to compute diffusion coefficient of replicas
        # Total exchange steps performed
        beta_stats[-1] += 1

        delta = beta_stats[0] - beta_stats[1]
        beta_stats[1] += delta / float(beta_stats[-1])
        delta2 = beta_stats[0] - beta_stats[1]
        beta_stats[2] += delta * delta2

    @property
    def stats(self):
        stats = {}
        accept = self._accepted_samples / float(self._total_samples)

        stats["mean_acceptance"] = _mean(accept)
        stats["min_acceptance"] = _mean(accept.min())
        stats["max_acceptance"] = _mean(accept.max())

        # Average position of beta=1
        # This is normalized and centered around zero
        # In the ideal case the average should be zero
        stats["normalized_beta=1_position"] = (
            self._beta_stats[1] / float(self._n_replicas - 1) - 0.5
        )

        # Average variance on the position of beta=1
        # In the ideal case this quantity should be of order ~ [0.2, 1]
        stats["normalized_beta=1_diffusion"] = _np.sqrt(
            self._beta_stats[2] / self._beta_stats[-1]
        ) / float(self._n_replicas)

        return stats

ext: py

max_stars_repo_name: gpescia/MyNetKet

max_issues_repo_name: gpescia/MyNetKet

max_forks_repo_name: gpescia/MyNetKet

reason: import numpy,from numba


text:
import numpy as np
import torch
from implem.utils import init_torch_device, as_tensor, device, dtype, dtype_np
from implem.utils import banded_gauss_seidel_redblack
import matplotlib.pyplot as plt

###########################################
# Semi-implicit numerical solvers for SWE #
###########################################

def swe1D_true_bM(x, dt, dx, g, w_imp, cd, data_scales, comp_u='solve'):
    # 1-D Shallow-Water equations semi-implicit solver
    # ported to python from Fortran code written by
    #               Kai Logemann, HZG, October 2020

    assert comp_u in ['solve', 'calculate']
    N = x.shape[0]

    ksi = x[:,0] * data_scales.flatten()[0]
    u   = x[:,1] * data_scales.flatten()[1]
    dp  = x[:,-2] * data_scales.flatten()[2]

    h = ksi + dp
    nx = ksi.shape[-1]

    Ce, Cw = torch.zeros_like(ksi), torch.zeros_like(ksi)
    uint, us, uints = torch.zeros_like(ksi), torch.zeros_like(ksi), torch.zeros_like(ksi)

    i, ip = np.arange(1, nx-2), np.arange(2, nx-1)
    hm = (h[:,i]+h[:,ip]) / 2.
    Cw[:,ip] = Ce[:,i] = dt**2/dx**2 * w_imp**2 * g * hm
    uint[:,i] = hm * u[:,i]
    us[:,i] = u[:,i] - dt*cd/hm*u[:,i]*torch.abs(u[:,i]) - (1.-w_imp)*dt*g*(ksi[:,ip]-ksi[:,i])/dx
    uints[:,i] = hm * us[:,i]

    if comp_u == 'solve':
        M_true = torch.zeros((N,2,2,3,1,nx))
    elif comp_u == 'calculate':
        M_true = torch.zeros((N,1,1,3,1,nx))

    # dksi_new/dksi_old
    M_true[:,0,0,0,0] = - Cw / (1. + Ce + Cw) # Cw: first two elements and last element zero
    M_true[:,0,0,1,0] = 1.
    M_true[:,0,0,2,0] = - Ce / (1. + Ce + Cw) # Ce: first element and last two elements zero

    if comp_u == 'solve':
        # will return system of linear equations to solve for both new xi and u

        # du_new/dksi_old
        M_true[:,1,0,1,0,i] = - dt/dx * w_imp * g # main diagonal and first upper diaginal:
        M_true[:,1,0,2,0,i] = dt/dx * w_imp * g   # first element and last two elements zero

        # du_new/du_old
        M_true[:,1,1,1,0] = 1.

        b_true = torch.zeros((N,2*nx))
        i, im = np.arange(1, nx-1), np.arange(0, nx-2)

        # water height
        # old state
        b_true[:,i] = ksi[:,i] # part of b for xi: first element and last element zero
        # div
        b_true[:,i] -= dt*(1.-w_imp)/dx * (uint[:,i] - uint[:,im])
        b_true[:,i] -= dt*w_imp/dx * (uints[:,i] - uints[:,im])
        # normalize rows of A & b:
        b_true[:,:nx] /= (1 + Ce + Cw)

        # velocities
        b_true[:,nx:] = us  # part of b for u: first element and last two elements zero  (via us)

        b_true = b_true.reshape(N,2,1,nx)

    if comp_u == 'calculate':
        # will return system of linear equations to solve only for new xi.
        # New u will be calculated from that in an extra step.

        b_true = torch.zeros((N,nx))
        i, im = np.arange(1, nx-1), np.arange(0, nx-2)

        # water height
        # old state
        b_true[:,:nx] = ksi
        # closed boundaries
        b_true[:,0] = b_true[:,nx-1] = 0.
        # div
        b_true[:,i] -= dt*(1.-w_imp)/dx * (uint[:,i] - uint[:,im])
        b_true[:,i] -= dt*w_imp/dx * (uints[:,i] - uints[:,im])
        # normalize rows of A & b:
        b_true[:,:nx] /= (1 + Ce + Cw)

        b_true = b_true.reshape(N,1,1,nx)

    return b_true, M_true, us

def swe1D_true_solve(x, dt, dx, g, w_imp, cd, data_scales, settings, comp_u='solve'):

    assert comp_u in ['solve', 'calculate']

    b_true, M_true, us = swe1D_true_bM(x, dt, dx, g, w_imp, cd, data_scales, comp_u)

    pred_hat, diagnostics = banded_gauss_seidel_redblack(M_true, b_true, **settings)

    if comp_u == 'solve':
        pred_hat_out = pred_hat

    elif comp_u == 'calculate':
        pred_hat_out = torch.zeros((x.shape[0],2,1,x.shape[-1]))
        pred_hat_out[:,:1] = pred_hat
        pred_hat_out[:,1,0,1:-2] = us[:,1:-2] - dt/dx * w_imp * g * (pred_hat_out[:,0,0,2:-1] - pred_hat_out[:,0,0,1:-2])

    return pred_hat_out, diagnostics

def swe1D_get_init(N, nx):

    # set depth profile d [m]
    sigma, tau, mu = 10., 100., 15
    profile = sigma * np.exp(- (np.arange(2*nx)- nx)**2 / tau)
    mean, cov = np.ones(nx) * mu, np.empty((nx,nx))
    cov[0] = profile[nx:]
    for i in range(1,nx):
        cov[i] = profile[nx-(i):-i]
    d = np.atleast_2d(np.random.multivariate_normal(mean, cov, size=N))
    d[:,0] = d[:,-1] = - mu
    assert np.all(d[:,1:-1] > 0) # no dry cells within domain for now

    # set initial sea surface elevation zeta [m]
    zeta = 0 * np.sqrt(0.001) * np.random.normal(size=(N,nx))
    zeta[np.arange(N),np.random.randint(nx, size=N)] = 0.05 + np.random.random(size=N) * 0.1 # initial disturbance
    # set initial velocities u [m/s]
    u = 0. * np.random.normal(size=(N,nx))

    mask = np.ones((N,nx))
    mask[:,0] = mask[:,-1] = 0.

    return zeta, u, d, mask

def swe1D_sim(N, T, nx, dt, dx, g, w_imp, cd, data_scales, settings,
              verbose=False, comp_u='solve', init_vals=None):

    zeta, u, d, mask = swe1D_get_init(N=N, nx=nx)
    x0 = np.stack((zeta, u, d, mask), axis=1)
    if not init_vals is None:
        assert init_vals.shape == x0.shape
        x0 = init_vals
        d, mask = init_vals[:,-2], init_vals[:,-1]

    statics = as_tensor(np.stack((d, mask), axis=1))

    out = [as_tensor(x0)]
    with torch.no_grad():
        for t in range(T):
            if verbose and np.mod(t,10)==0:
                print(f'simulating {t}/{T}')
            #x_est = model(x)
            x_est, diagnostics = swe1D_true_solve(out[-1], dt, dx, g, w_imp, cd,
                                             data_scales, settings=settings, comp_u=comp_u)
            x_est = x_est.squeeze(-2)/as_tensor(data_scales[0,:,:2])
            out.append(torch.cat((x_est, statics), dim=1))

    return torch.stack(out, axis=1)

def swe2D_true_bM(x, dt, dx, g, w_imp, cd, ah, data_scales, comp_u='solve'):
    # 2-D Shallow-Water equations semi-implicit solver
    # ported to python from Fortran code written by
    #               Kai Logemann, HZG, October 2020

    N = x.shape[0]
    ny, nx = x.shape[-2:]

    ksi = x[:,0] * data_scales.flatten()[0]
    u   = x[:,1] * data_scales.flatten()[1]
    v   = x[:,2] * data_scales.flatten()[2]
    dp  = x[:,3] * data_scales.flatten()[3]

    h = ksi + dp

    # equation of motion => us,vs (interim solution)
    # vertical integrals of velocity => uint, and interim solution uints
    # coefficients for zetan system of equations: ce,cw,cn,cs

    Ce, Cw = torch.zeros_like(ksi), torch.zeros_like(ksi)
    Cn, Cs = torch.zeros_like(ksi), torch.zeros_like(ksi)
    uint, us, uints = torch.zeros_like(ksi), torch.zeros_like(ksi), torch.zeros_like(ksi)

    i, ip, im = np.arange(1, ny-2), np.arange(2, ny-1), np.arange(0, ny-3)
    j, jp, jm = np.arange(1, nx-1), np.arange(2, nx), np.arange(0, nx-2)
    hm = (h[:,i][:,:,j]+h[:,ip][:,:,j]) / 2.
    uij = u[:,i][:,:,j]
    usij = uij - dt*cd/hm*uij*torch.abs(uij)
    usij -= (1.-w_imp)*dt*g*(ksi[:,ip][:,:,j]-ksi[:,i][:,:,j])/dx
    usij += dt*ah*(u[:,ip][:,:,j]+ u[:,im][:,:,j] + u[:,i][:,:,jp] + u[:,i][:,:,jm] - 4.*uij)/dx**2
    usi, uinti, uintsi = us[:,i], uint[:,i], uints[:,i]
    usi[:,:,j] = usij
    us[:,i] = usi
    uinti[:,:,j] = hm * uij
    uint[:,i] = uinti
    uintsi[:,:,j] = hm * usij
    uints[:,i] = uintsi

    Cei = Ce[:,i]
    Cei[:,:,j] = dt**2/dx**2 * w_imp**2 * g * hm
    Ce[:,i] = Cei

    vint, vs, vints = torch.zeros_like(ksi), torch.zeros_like(ksi), torch.zeros_like(ksi)
    i, ip, im = np.arange(1, ny-1), np.arange(2, ny), np.arange(0, ny-2)
    j, jp, jm = np.arange(1, nx-2), np.arange(2, nx-1), np.arange(0, nx-3)
    hm = (h[:,i][:,:,j]+h[:,i][:,:,jp]) / 2.
    vij = v[:,i][:,:,j]
    vsij = vij - dt*cd/hm*vij*torch.abs(vij)
    vsij -= (1.-w_imp)*dt*g*(ksi[:,i][:,:,jp]-ksi[:,i][:,:,j])/dx
    vsij += dt*ah*(v[:,ip][:,:,j]+ v[:,im][:,:,j] + v[:,i][:,:,jp] + v[:,i][:,:,jm] - 4.*vij)/dx**2

    vsi, vinti, vintsi = vs[:,i], vint[:,i], vints[:,i]
    vsi[:,:,j] = vsij
    vs[:,i] = vsi
    vinti[:,:,j] = hm * vij
    vint[:,i] = vinti
    vintsi[:,:,j] = hm * vsij
    vints[:,i] = vintsi

    Cni = Cn[:,i]
    Cni[:,:,j] = dt**2/dx**2 * w_imp**2 * g * hm
    Cn[:,i] = Cni

    i, im, j = np.arange(2, ny-1), np.arange(1, ny-2), np.arange(1, nx-1)
    hm = (h[:,i][:,:,j]+h[:,im][:,:,j]) / 2.
    Cwi = Cw[:,i]
    Cwi[:,:,j] = dt**2/dx**2 * w_imp**2 * g * hm
    Cw[:,i] = Cwi

    i, j, jm = np.arange(1, ny-1), np.arange(2, nx-1), np.arange(1, nx-2)
    hm = (h[:,i][:,:,j]+h[:,i][:,:,jm]) / 2.
    Csi = Cs[:,i]
    Csi[:,:,j] = dt**2/dx**2 * w_imp**2 * g * hm
    Cs[:,i] = Csi

    if comp_u == 'solve':
        M_true = torch.zeros((N,3,3,5,ny,nx))
    elif comp_u == 'calculate':
        M_true = torch.zeros((N,1,5,ny,nx))

    M_true[:,0,0] = - Cw / (1. + Ce + Cw + Cs + Cn) # in the Fortran code, i+/-1 signifies east-west
    M_true[:,0,1] = - Cs / (1. + Ce + Cw + Cs + Cn) #                      j+/-1 signifies north-south
    M_true[:,0,2] = 1.                              # torch.nn.unfold works row-wise (first row, second row...)
    M_true[:,0,3] = - Cn / (1. + Ce + Cw + Cs + Cn) # i.e. A*vec(x)=b means vec(x) is first row, second row...
    M_true[:,0,4] = - Ce / (1. + Ce + Cw + Cs + Cn) # and hence [[     0=W     ]
                                                    #            [1=S   2   3=N] for a 3x3 stencil ... ?
                                                    #            [     4=E   0]]

    if comp_u == 'solve':
        raise NotImplementedError()

    elif comp_u == 'calculate':

        b_true = torch.zeros((N, ny, nx))
        i, im = np.arange(1, ny-1), np.arange(0, ny-2)
        j, jm = np.arange(1, nx-1), np.arange(0, nx-2)

        # water height
        # old state
        b_true[:,:ny] = ksi
        # closed boundaries
        b_true[:,0] = b_true[:,ny-1] = 0.
        b_true[:,:,0] = b_true[:,:,nx-1] = 0.
        # div
        b_truei = b_true[:,i]
        b_truei[:,:,j] -= dt*(1.-w_imp)/dx * (uint[:,i][:,:,j] - uint[:,im][:,:,j])
        b_truei[:,:,j] -= dt*w_imp/dx * (uints[:,i][:,:,j] - uints[:,im][:,:,j])
        b_truei[:,:,j] -= dt*(1.-w_imp)/dx * (vint[:,i][:,:,j] - vint[:,i][:,:,jm])
        b_truei[:,:,j] -= dt*w_imp/dx * (vints[:,i][:,:,j] - vints[:,i][:,:,jm])
        b_true[:,i] = b_truei

        # normalize rows of A & b:
        b_true[:,:ny] /= (1. + Ce + Cw + Cn + Cs)

        b_true = b_true.reshape(N,1,ny,nx)

    return b_true, M_true, us, vs

def swe2D_true_solve(x, dt, dx, g, w_imp, cd, ah, data_scales, settings, comp_u='solve'):

    assert comp_u in ['solve', 'calculate']
    assert settings['stencil_offsets'].shape[1] == 5 # 5-point stencil !

    b_true, M_true, us, vs = swe2D_true_bM(x, dt, dx, g, w_imp, cd, ah, data_scales, comp_u)

    pred_hat, diagnostics = banded_gauss_seidel_redblack(M_true, b_true, **settings)

    if comp_u == 'solve':
        raise NotImplementedError()

    elif comp_u == 'calculate':

        zetan = pred_hat[:,0]
        ny, nx = zetan.shape[-2:]

        un, vn = torch.zeros_like(us), torch.zeros_like(vs)
        i, ip, j  = np.arange(1, ny-2), np.arange(2, ny-1), np.arange(1, nx-1)

        uni = un[:,i]
        uni[:,:,j] = us[:,i][:,:,j] - w_imp*dt/dx*g*(zetan[:,ip][:,:,j]-zetan[:,i][:,:,j])
        un[:,i] = uni

        i, j, jp = np.arange(1, ny-1), np.arange(1, nx-2), np.arange(2, nx-1)

        vni = vn[:,i]
        vni[:,:,j] = vs[:,i][:,:,j] - w_imp*dt/dx*g*(zetan[:,i][:,:,jp]-zetan[:,i][:,:,j])
        vn[:,i] = vni

        pred_hat_out = torch.stack((zetan, un, vn), axis=1)

    return pred_hat_out, diagnostics

def swe2D_get_init(N, my, nx):

    # set constant depth profile d [m]
    sigma, tau, mu = 5., 100., 15
    mean = np.ones((nx*my,1)) * mu
    x,y = np.stack(np.meshgrid(np.arange(my), np.arange(nx))).reshape(2,-1)
    dists = (x.reshape(-1,1) - x.reshape(1,-1))**2 + (y.reshape(-1,1) - y.reshape(1,-1))**2
    cov = sigma * np.exp(- dists / tau) + 1e-4 * np.eye(nx*my)
    L = np.linalg.cholesky(cov)
    d = np.random.normal(size=(nx*my,N))
    d = (L.dot(d) + mean).reshape(nx,my,N)
    d[:,0] = d[:,-1] = -mu
    d[0,:] = d[-1,:] = - mu
    assert np.all(d[1:-1][:,1:-1] > 0) # no dry cells within domain for now

    # set initial sea surface elevation zeta [m]
    zeta = 0. * np.sqrt(0.00001) * np.random.normal(size=(nx,my,N))

    for i in range(N):
        idn, idm = np.random.randint(nx), np.random.randint(my)
        pert_size = 5 + np.random.randint(6)
        idn = range(max(0, -(pert_size//2)+idn),  min(nx, (pert_size+1)//2+idn))
        idm = range(max(0, -(pert_size//2)+idm),  min(my, (pert_size+1)//2+idm))
        zetai = zeta[:,:,i]
        zetai[np.ix_(idn, idm)] = 0.05 + np.random.random() * 0.1 # initial disturbance
        zeta[:,:,i] = zetai

    # set initial velocities u [m/s]
    u = 0. * np.random.normal(size=(nx,my,N))
    v = 0. * np.random.normal(size=(nx,my,N))

    # boundary mask (boundaries = dry cells)
    mask = np.ones((N,my,nx))
    mask[:,0] = mask[:,-1] = 0.
    mask[:,:,0] = mask[:,:,-1] = 0.

    return zeta.transpose(), u.transpose(), v.transpose(), d.transpose(), mask

def swe2D_sim(N, T, my, nx, dt, dx, g, w_imp, cd, ah, data_scales, settings,
              verbose=False, comp_u='calculate', init_vals=None):

    zeta, u, v, d, mask = swe2D_get_init(N=N, nx=nx, my=my)
    if not init_vals is None:
        assert init_vals.shape == d.shape
        d = 1. * init_vals
    statics = as_tensor(np.stack((d, mask), axis=1))

    out = [as_tensor(np.stack((zeta, u, v, d + zeta, mask), axis=1))] # d+zeta because 2D code expects h!
    with torch.no_grad():
        for t in range(T):
            if verbose and np.mod(t,10)==0:
                print(f'simulating {t}/{T}')
            x_est, diagnostics = swe2D_true_solve(out[-1], dt, dx, g, w_imp, cd, ah,
                                             data_scales, settings=settings, comp_u=comp_u)
            x_est = x_est.squeeze(-2)/as_tensor(data_scales[0,:,:3])
            out.append(torch.cat((x_est, statics), dim=1))
            out[-2] = out[-2].detach().cpu() # freeing up GPU memory
    out[-1] = out[-1].detach().cpu()

    return torch.stack(out, axis=1)

def swe_sim(N, T, setup, solver_settings, verbose=False):

    assert setup['instance_dimensionality'] in ['1D', '2D']
    if not 'w_imp' in setup.keys():
        setup['w_imp'] = 0.5

    if setup['instance_dimensionality'] == '1D':
        if not 'data_scales' in setup.keys():
                setup['data_scales'] = np.ones((1,1,4,1))
        out = swe1D_sim(N, T,
                        nx=setup['nx'],
                        dt=setup['dt'],
                        dx=setup['dx'],
                        g=setup['g'],
                        w_imp=setup['w_imp'],
                        cd=setup['cd'],
                        data_scales=setup['data_scales'],
                        settings=solver_settings,
                        verbose=verbose,
                        comp_u=setup['comp_u'])
    else:
        if not 'data_scales' in setup.keys():
            setup['data_scales'] = np.ones((1,1,5,1,1))
        if not 'my' in setup.keys():
            setup['my'] = 1 * setup['nx']
        out = swe2D_sim(N, T,
                        my=setup['my'],
                        nx=setup['nx'],
                        dt=setup['dt'],
                        dx=setup['dx'],
                        g=setup['g'],
                        w_imp=setup['w_imp'],
                        cd=setup['cd'],
                        ah=setup['ah'],
                        data_scales=setup['data_scales'],
                        settings=solver_settings,
                        verbose=verbose,
                        comp_u=setup['comp_u'])
    return out

def plot_results_swe(data_numerical, data_model, i=0, swe_model='1D', if_save=False, fig_path=None, plot_times=None):

    if if_save:
        assert not fig_path is None

    assert swe_model in ['1D', '2D']

    fontsize = 14
    if swe_model == '1D':
        plt.figure(figsize=(16,10))
        x, x_hat = data_numerical, data_model

        vmax_diff = np.abs(x - x_hat).max()
        vmin, vmax = np.min((x.min(), x_hat.min())), np.max((x.max(), x_hat.max()))
        plt.subplot(3,1,1)
        plt.imshow(x.T, vmin=vmin, vmax=vmax, cmap='viridis', aspect='auto')
        plt.xticks([])
        plt.yticks([])
        plt.ylabel('data', fontsize=fontsize)
        plt.colorbar()

        plt.subplot(3,1,2)
        plt.imshow(x_hat.T, vmin=vmin, vmax=vmax, cmap='viridis', aspect='auto')
        plt.xticks([])
        plt.yticks([])
        plt.ylabel('model', fontsize=fontsize)
        plt.colorbar()

        plt.subplot(3,1,3)
        plt.imshow(x.T - x_hat.T, vmin=-vmax_diff, vmax=vmax_diff, cmap='bwr', aspect='auto')
        plt.yticks([])
        plt.xticks(fontsize=fontsize)
        plt.ylabel('difference', fontsize=fontsize)
        plt.xlabel('time step', fontsize=fontsize)
        plt.colorbar()

        if i == 0:
            plt.suptitle('water depth', fontsize=fontsize, x=0.45, y=0.93)
        elif i == 1:
            plt.suptitle('velocity u', fontsize=fontsize, x=0.45,y=0.93)
        elif i == 2:
            plt.suptitle('mask', fontsize=fontsize, x=0.45, y=0.93)

    elif swe_model == '2D':

        plot_times = np.arange(0, len(data_model), len(data_model)//4) if plot_times is None else plot_times
        num_times = len(plot_times)
        plt.figure(figsize=(16,14))
        x, x_hat = data_numerical, data_model

        vmax_diff = np.abs(x[plot_times] - x_hat[plot_times]).max()

        for j,t in enumerate(plot_times):
            vmin, vmax = np.min((x[t].min(), x_hat[t].min())), np.max((x[t].max(), x_hat[t].max()))
            plt.subplot(len(plot_times),num_times,1+num_times+j)
            plt.imshow(x[t].T, vmin=vmin, vmax=vmax, cmap='viridis', aspect='auto')
            plt.xticks([])
            plt.yticks([])
            plt.title(f't = {t}')
            if j == 0:
                plt.ylabel('data', fontsize=fontsize)
            plt.colorbar()
            plt.subplot(len(plot_times),num_times,1+(2)*num_times+j)
            plt.imshow(x_hat[t].T, vmin=vmin, vmax=vmax, cmap='viridis', aspect='auto')
            plt.xticks([])
            plt.yticks([])
            if j == 0:
                plt.ylabel('model', fontsize=fontsize)
            plt.colorbar()
            plt.subplot(len(plot_times),num_times,1+(3)*num_times+j)
            plt.imshow(x[t].T - x_hat[t].T, vmin=-vmax_diff, vmax=vmax_diff, cmap='bwr', aspect='auto')
            plt.xticks([])
            plt.yticks([])
            if j == 0:
                plt.ylabel('difference', fontsize=fontsize)
            plt.colorbar()

        if i == 0:
            plt.suptitle('water depth', fontsize=fontsize, x=0.45, y=0.73)
        elif i == 1:
            plt.suptitle('velocity u', fontsize=fontsize, x=0.45,y=0.73)
        elif i == 2:
            plt.suptitle('velocity v', fontsize=fontsize, x=0.45,y=0.73)

    if if_save:
        if i == 0:
            plt.savefig(fig_path + 'example_rollout_water_depth_2D_no_boundaries.pdf')
        elif i == 1:
            plt.savefig(fig_path + 'example_rollout_velocities_u_2D_no_boundaries.pdf')
        elif i == 2:
            plt.savefig(fig_path + 'example_rollout_velocities_v_2D_no_boundaries.pdf')

    plt.show()

################################################
# SWE-specific code for implem implicit models #
################################################

def enforce_bcs(x, z):
    # boundary conditions right on boundary pixels. x = (water height, u, v) is full state !
    cond_mask = z[:,-1:] # z could be data point, has to have boundary mask on last channel.
    x[:,:3] = cond_mask * x[:,:3]
    x[:,1,-2,:] = 0. # u at southern border is zero
    x[:,2,:,-2] = 0. # v at eastern border is zero
    return x

def fix_bM_BCs(settings, x, b, M):
    # ensures that boundary conditions are met by externally overwriting parts of M,b.
    # Note that when writing M = mask * Md + (1-mask) * M in torch, gradients are tracked!

    cond_mask = 1. - x[:,-1:]
    with torch.no_grad():
        ba = b[:,:1].detach()
        Ma = M[:,:1].detach() if settings['sep_eqs_per_field'] else M[:,0,:1].detach()
        bd, Md = gen_bM_fixed_BCs(settings=settings,
                                cond_mask=cond_mask,
                                M=Ma, b=ba)

    # paste in manualy written BC equations over boundary grid points:
    if settings['sep_eqs_per_field']:
        M[:,:1] = cond_mask.unsqueeze(-3) * Md + (1.-cond_mask.unsqueeze(-3)) * M[:,:1]
    else:
        # M[:,0,:1] dependence of first field on itself, M[:,0,:1] dependence on other fields.
        M[:,0,:1] = cond_mask.unsqueeze(-3) * Md + (1.-cond_mask.unsqueeze(-3)) * M[:,0,:1]
        M[:,0,1:] = cond_mask.unsqueeze(-3) * 0. + (1.-cond_mask.unsqueeze(-3)) * M[:,0,1:]
    b[:,:1] = cond_mask * bd + (1.-cond_mask) * b[:,:1]

    # tricky one: boundary conditions bleed into neighbouring pixels !
    # could be more elegantly done by convolving the boundary mask cond_mask with
    # kernels [0,0,1], [1,0,0] resp. [[0],[0],[1]], [[1],[0],[0]] to find grid points
    # neighouring the respective boundary
    stencil_offsets = settings['stencil_offsets'].detach().cpu().numpy()
    idxt = np.where((stencil_offsets[0] == -1) * (stencil_offsets[1] == 0))[0][0] # top
    M[...,idxt, 1, :] =  0. # Northern end
    idxb = np.where((stencil_offsets[0] ==  1) * (stencil_offsets[1] == 0))[0][0] # bottom
    M[...,idxb,-2, :] =  0. # Southern end
    idxl = np.where((stencil_offsets[0] ==  0) * (stencil_offsets[1] ==-1))[0][0] # left
    M[...,idxl, :, 1] =  0. # Western end
    idxr = np.where((stencil_offsets[0] ==  0) * (stencil_offsets[1] == 1))[0][0] # right
    M[...,idxr, :,-2] =  0. # Eastern end

    backwards_mask = None

    return b, M, backwards_mask

def gen_bM_fixed_BCs(settings, cond_mask, b=None, M=None):
    # ensure that the system of linear equations defined by M,b have equations
    # 1. * x_i = 0.
    # for boundary grid points x_i, meaning we force these values to be zero.
    # Note that this is different from just zero-ing these values after the solve,
    # since the grid points within the domain could depend on them during the solve!

    ny, nx = cond_mask.shape[-2:]
    stencil_offsets = settings['stencil_offsets'].detach().cpu().numpy()

    # output matrix for system of linear equations that satisfies BCs
    M_shape = (cond_mask.shape[0],1,stencil_offsets.shape[1],ny,nx)
    if M is None:
        M = torch.zeros(M_shape,dtype=v_cond.dtype)
        M[:,:,idxc] = 1.
    else:
        assert np.all(M.shape == M_shape)
        M = 0. * cond_mask.unsqueeze(-3) + (1 - cond_mask.unsqueeze(-3)) * M

    # retrieve axis indices for different stencil positions (we need a 5-point cross stencil)
    idxc = np.where((stencil_offsets[0] ==  0) * (stencil_offsets[1] == 0))[0][0] # center
    M[:,:,idxc] =  1. * cond_mask + (1 - cond_mask) * M[:,:,idxc] # center

    # right-hand side of system of linear equations
    b_shape = (cond_mask.shape[0],1,ny,nx)
    if b is None:
        b = torch.zeros(b_shape,dtype=v_cond.dtype)
    else:
        assert np.all(b.shape == b_shape)

    b[...,0] = b[...,-1] = 0.
    b[...,0,:] = b[...,-1,:] = 0.

    return b, M

def network_init_swe(model, x, data_scales, dt=300., dx=1e4, w_imp=0.5, g=9.81, cd=1e-3, ah=1e3):
    # network initialization for implicit network to solve 2D SWE
    # ensures that water height is computed in the first implicit layer (one or zero solved fields)
    # ensures that u,v are computed in the second implicit layer with zero solved fields (explicit mode!)

    # first implicit layer (block of several conv layers!)
    if model.impl_layers[0].hparams.implicit_mode: # if training (semi-)implicit model

        # initialize deep implicit layer to some constant fixed 5-point stencil
        dp  = x[:,3] * data_scales.flatten()[3]
        fcs = dp.mean() * dt**2/dx**2 * w_imp**2 * g # typical values for Cn,Cs,Cw,Ce averaged across space
        final_forward_layer = model.impl_layers[0].forward_model.layers[-1]
        bias, weight = final_forward_layer.bias, final_forward_layer.weight
        bias_new, weight_new = 1.* bias, 1.* weight
        bias_new[:5] = - fcs + 0.*bias[:5]
        bias_new[2] = 1. + 0.*bias[2]
        weight_new[:5] = weight_new[:5]/1e2
        final_forward_layer.bias = torch.nn.Parameter(bias_new)
        final_forward_layer.weight = torch.nn.Parameter(weight_new)

        # 50 iterations are okay for convergence of systems for dt=300,900,1500, but beyond
        # the diagonal dominance of the true matrix A is so weak that more iterations are needed.
        model.impl_layers[0].settings['max_iter_forward'] = 50
        model.impl_layers[0].settings['max_iter_backward'] = 50
        model.impl_layers[0].settings['thresh_forward'] = 1e-14  # input normalized to scale ~1
        model.impl_layers[0].settings['thresh_backward'] = 1e-24 # higher tolerance since dL/dz can be tiny!

    if not model.impl_layers[0].hparams.implicit_mode:  # if training purely explicit model
        pass

    # fix first out_conv to ensure passing on results of implicit layer output to second out_conv (skip connection!)
    n = model.impl_layers[0].out_conv.weight.shape[1]
    model.impl_layers[0].out_conv.weight = torch.nn.Parameter(torch.eye(n,dtype=x.dtype).unsqueeze(-1).unsqueeze(-1))
    model.impl_layers[0].out_conv.bias = torch.nn.Parameter(torch.zeros_like(model.impl_layers[0].out_conv.bias))
    model.impl_layers[0].out_conv.weight.requires_grad = False
    model.impl_layers[0].out_conv.bias.requires_grad = False

    # second `implicit` layer: here explicit layer with single 3x3 conv !

    # fix second out_conv to copy water height from first layer, and u,v from second layer
    # we fix this also for explicit models for fair comparison
    # (otherwise they have 1 more 3x3 conv for water height)!
    n = model.impl_layers[1].out_conv.weight.shape[1]
    mat = torch.zeros_like(model.impl_layers[1].out_conv.weight)
    mat[1,0] = 1. # ordering of input is: (out [0 channels], out_expl [2 channels], skip connection [5+ channels]
    mat[2,1] = 1. # so incoming indeces 0,1 take input from out_expl of second layer,
    mat[0,2] = 1. # and index 2 takes input from first channel of skip connection from first out_conv
    model.impl_layers[1].out_conv.weight = torch.nn.Parameter(mat)
    model.impl_layers[1].out_conv.bias = torch.nn.Parameter(torch.zeros_like(model.impl_layers[1].out_conv.bias))
    model.impl_layers[1].out_conv.weight.requires_grad = False
    model.impl_layers[1].out_conv.bias.requires_grad = False

reason: import numpy


text:
import os
import numpy as np
import torch

from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

import cv2
import random
import math

import pytorch_lightning as pl
from torch.utils.data import DataLoader

from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import random_split, DataLoader

class SprayDataModule(pl.LightningDataModule):

    def __init__(self, dataset_parms, trainer_parms, model_params, data_dir = './'):
        super().__init__()
        self.dataset_parms = dataset_parms
        self.trainer_parms = trainer_parms

        self.dims = (model_params['in_channels'], dataset_parms['img_size'], dataset_parms['img_size'])
        #self._has_prepared_data = True
        #self._has_setup_fit = True

        self.data_dir = data_dir
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])

    def prepare_data(self):
        print("prepare_data")
        # download
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):
        '''
        if stage == 'fit' or stage is None:
            self.data_train = SprayDataset(self.dataset_parms['train_list'],
                                           crop_size=self.dataset_parms['img_size'])

            self.data_val = SprayDataset(self.dataset_parms['val_list'],
                                         crop_size=self.dataset_parms['img_size'])
        '''
         # Assign train/val datasets for use in dataloaders
        if stage == 'fit' or stage is None:
            mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
            self.data_train, self.data_val = random_split(mnist_full, [55000, 5000])

            # Optionally...
            # self.dims = tuple(self.mnist_train[0][0].shape)

        # Assign test dataset for use in dataloader(s)
        if stage == 'test' or stage is None:
            self.data_test = MNIST(self.data_dir, train=False, transform=self.transform)

            # Optionally...
            # self.dims = tuple(self.mnist_test[0][0].shape)

    def train_dataloader(self):
        return DataLoader(self.data_train,
                          batch_size=self.dataset_parms['batch_size'],
                          shuffle=True,
                          num_workers=self.trainer_parms['n_workers'])

    def val_dataloader(self):
        return DataLoader(self.data_val,
                          batch_size=self.dataset_parms['batch_size'],
                          shuffle=False,
                          num_workers=self.trainer_parms['n_workers'],
                          pin_memory=True)

    def test_dataloader(self):
        return DataLoader(self.data_test,
                          batch_size=self.dataset_parms['batch_size'],
                          shuffle=False,
                          num_workers=self.trainer_parms['n_workers'],
                          pin_memory=True)

class SprayDataset(torch.utils.data.Dataset):
    def __init__(self, list_path, crop_size=128):
        self.crop_size = crop_size
        self.transforms = self.get_transform()

        with open(list_path) as f:
            self.label_list = f.read().splitlines()

    def load_sample(self, label_path):

        with open(label_path) as f:
            line = f.read().splitlines()[0]
            values = []
            #point_t, point_l, point_r, line_l, line_r
            for set in line.rstrip().split(';'):
                # x,y or a,b
                for str_val in set.split(':')[1].split(','):
                    values.append(float(str_val))
        label = values[:]

        image_path = label_path.replace('txt','jpg')
        #img = Image.open(self.img_list[idx]).convert("RGB")
        img = cv2.imread(image_path)
        self.image_h, self.image_w, _ = img.shape
        img = cv2.cvtColor(img[:self.image_h,:self.image_h], cv2.COLOR_BGR2RGB)

        return img, label

    def get_transform(self):
        tfms = []
        tfms.append(transforms.ToTensor())
        #tfms.append(Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225]))
        return transforms.Compose(tfms)

    def rot_and_crop(self, img, values, show=False):
        image_h, image_w, _ = img.shape
        cX, cY = image_w//2, image_h//2

        angle = random.randint(0,180)
        M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)

        # compute the new bounding dimensions of the image
        #cos = np.abs(M[0, 0])
        #sin = np.abs(M[0, 1])
        #nW = int((image_h * sin) + (image_w * cos))
        #nH = int((image_h * cos) + (image_w * sin))

        # adjust the rotation matrix to take into account translation
        #M[0, 2] += (nW / 2) - cX
        #M[1, 2] += (nH / 2) - cY

        # original
        top_x, top_y = int(values[0]*self.image_w), int(values[1]*self.image_h)
        if show:
            img = cv2.circle(img, (top_x, top_y), 2, (0,255,255), 5)

        # rotate image
        img = cv2.warpAffine(img, M, img.shape[:2])

        # rotate points
        # method 1
        top = M.dot(np.array((top_x, top_y, 1)))
        top_x, top_y = int(top[0]), int(top[1])
        # method 2
        #radian = math.radians(angle)
        #top_x_r = int(cX + math.cos(radian) * (top_x - cX) - math.sin(radian) * (top_y - cY))
        #top_y_r = int(cY + math.sin(radian) * (top_x - cX) + math.cos(radian) * (top_y - cY))

        # crop around new rotated top point
        crop_offset = 20
        x = max(0, top_x - self.crop_size//2 + random.randint(-crop_offset,crop_offset))
        y = max(0, top_y - self.crop_size//2 + random.randint(-crop_offset,crop_offset))
        x = min(x, self.image_h - self.crop_size)
        y = min(y, self.image_h - self.crop_size)
        crop = img[y:y+self.crop_size, x:x+self.crop_size]

        if show:
            # draw new top point in crop
            crop = cv2.circle(crop, (top_x-x, top_y-y), 2, (0,0,255), 3)

        #left_x, left_y = int(values[2]*image_w), int(values[3]*image_h)
        #crop = cv2.circle(crop, (left_x-x, left_y-y), 2, (0,0,255), 3)

        #right_x, right_y = int(values[4]*image_w), int(values[5]*image_h)
        #crop = cv2.circle(crop, (right_x-x, right_y-y), 2, (0,0,255), 3)

        if show:
            cv2.imshow("crop", crop)
            cv2.waitKey()

        return crop

    def __getitem__(self, idx):
        img, target = self.load_sample(self.label_list[idx])
        img = self.rot_and_crop(img, target)

        #img = self.transforms(img)

        img = img.transpose((2, 0, 1))
        img = img / 255.0
        #img[0] = (img[0] - 0.485)/0.229
        #img[1] = (img[1] - 0.456)/0.224
        #img[2] = (img[2] - 0.406)/0.225
        img = torch.from_numpy(img)
        img = img.float()
        return img, img#, target

    def __len__(self):
        return len(self.label_list)

if __name__ == '__main__':
    dataset = SprayDataset('/home/markpp/datasets/teejet/iphone_data/train_list.txt')

    print(dataset[10])

max_stars_repo_name: markpp/autoencoders

max_issues_repo_name: markpp/autoencoders

max_forks_repo_name: markpp/autoencoders


text:
__author__ = 'Varun Nayyar'

import numpy as np
import os
import cPickle

emotions = ["Neutral", "Sad", "Bored", "Happy", "HotAnger"]
speakers = ["CC1", "CL1", "GG1", "JG1", "MF1", "MK1", "MM1"]

ExtractedMFCCs = "../MFCCData"

def main():
    MFCCFiles = os.listdir(ExtractedMFCCs)
    try:
        MFCCFiles.remove('.DS_Store')
    except ValueError:
        pass
        # It didn't need to be removed
    MFCCVals = []
    for file in MFCCFiles:
        if "Sad" in file:
            print file
            with open(os.path.join(ExtractedMFCCs, file)) as f:
                speakerEmotion = cPickle.load(f)
                speakerEmotion = np.vstack(speakerEmotion)
                MFCCVals.append(speakerEmotion)

    SadCorpus = np.vstack(MFCCVals)

    print SadCorpus.shape

def SadCorpus():
    """
    A list of emotions
    :return: list of len 7. speakers voice for the sad emotion
    """
    MFCCFiles = os.listdir(ExtractedMFCCs)
    try:
        MFCCFiles.remove('.DS_Store')
    except ValueError:
        pass
        # It didn't need to be removed
    MFCCVals = []
    for file in MFCCFiles:
        if "Sad" in file:
            # print file
            with open(os.path.join(ExtractedMFCCs, file)) as f:
                speakerEmotion = cPickle.load(f)
                speakerEmotion = np.vstack(speakerEmotion)
                MFCCVals.append(speakerEmotion)

    return np.vstack(MFCCVals)

def getIndiviudalData(emotion, speakerID):
    """
    Mostly for testing, so we will maintain lists
    This will be a list of that speakers utterances of that word
    :param emotion:
    :param speakerID:
    :return:
    """
    if emotion not in emotions or speakerID not in speakers:
        raise Exception("No Such speaker: {} or emotion: {}".format(speakerID, emotion))
    #errorCheck

    filename = emotion+"_"+speakerID
    with open(os.path.join(ExtractedMFCCs, filename)) as f:
        speakerEmotion = cPickle.load(f)

    #No Vstacking, getting the list of each utterance
    return speakerEmotion

def getDecimatedCorpus(emotion, speakerID):
    deciLoc = "../DecimatedMFCCs/MFCC{}-{}.txt".format(emotion, speakerID)
    # print deciLoc
    with open(deciLoc, 'r') as f:
        Xpoints = cPickle.load(f)

    return Xpoints

def pickleTestCorpus(datSize = 1000):
    # import cPickle
    for speakerID in speakers:
        for emotion in emotions:
            Xpoints = getFullCorpus(emotion, speakerID)
            indexes = np.random.choice(Xpoints.shape[0], size = datSize, replace=False)
            Xpoints = Xpoints[indexes]
            with open("../DecimatedMFCCs/MFCC{}-{}.txt".format(emotion, speakerID), 'w') as f:
                cPickle.dump(Xpoints, f)

def getFullCorpus(emotion, speakerID = None):
    """
    Return the 6 speakers files in a massive vstack
    :param emotion:
    :param speakerID:
    :return:
    """

    if emotion not in emotions or (speakerID is not None and speakerID not in speakers):
        raise Exception("No Such speaker: {} or emotion: {}".format(speakerID, emotion))
    #error check

    if speakerID is None:
        #return whole corpus
        speakerID = "Derpington"
        # should not be in file

    MFCCFiles = os.listdir(ExtractedMFCCs)

    try:
        MFCCFiles.remove('.DS_Store')
    except ValueError:
        pass
        # It didn't need to be removed
    MFCCVals = []
    for file in MFCCFiles:
        if emotion in file and speakerID not in file:
            # print "Currently reading", file
            with open(os.path.join(ExtractedMFCCs, file)) as f:
                speakerEmotion = cPickle.load(f)
                speakerEmotion = np.vstack(speakerEmotion)
                MFCCVals.append(speakerEmotion)

    return np.vstack(MFCCVals)

def getCorpus(emotion, speakerID):
    return getDecimatedCorpus(emotion, speakerID)

def myFunct(dat1 = 5, dat2 = 10):
    print dat1 - dat2

def test():
    myFunct(10,3)

if __name__ == "__main__":
    test()
    # pass

max_stars_repo_name: nayyarv/MonteGMM

max_issues_repo_name: nayyarv/MonteGMM

max_forks_repo_name: nayyarv/MonteGMM


text:
#!/usr/bin/python3

import argparse
import math
import numpy as np
import os.path
import pickle
import sys

sys.path.append('../lib')
import Groups
import ProjectMgr

import match_culling as cull

r2d = 180.0 / math.pi

parser = argparse.ArgumentParser(description='Keypoint projection.')
parser.add_argument('--project', required=True, help='project directory')
args = parser.parse_args()

proj = ProjectMgr.ProjectMgr(args.project)
proj.load_images_info()

#source = 'matches_direct'
source = 'matches_grouped'
print("Loading matches:", source)
matches_orig = pickle.load( open( os.path.join(args.project, source), "rb" ) )
print('Number of original features:', len(matches_orig))
print("Loading optimized matches: matches_opt")
matches_opt = pickle.load( open( os.path.join(args.project, "matches_opt"), "rb" ) )
print('Number of optimized features:', len(matches_opt))

# load the group connections within the image set
groups = Groups.load(args.project)
print('Main group size:', len(groups[0]))

pair_angles = []
for i in range(len(proj.image_list)):
    p = [ [] for j in range(len(proj.image_list)) ]
    pair_angles.append(p)

def compute_angle(ned1, ned2, ned3):
    vec1 = np.array(ned3) - np.array(ned1)
    vec2 = np.array(ned3) - np.array(ned2)
    n1 = np.linalg.norm(vec1)
    n2 = np.linalg.norm(vec2)
    denom = n1*n2
    if abs(denom - 0.000001) > 0:
        try:
            tmp = np.dot(vec1, vec2) / denom
            if tmp > 1.0: tmp = 1.0
            return math.acos(tmp)
        except:
            print('vec1:', vec1, 'vec2', vec2, 'dot:', np.dot(vec1, vec2))
            print('denom:', denom)
            return 0
    else:
        return 0

print("Computing match pair angles...")
for match in matches_opt:
    for m1 in match[1:]:
        for m2 in match[1:]:
            if m1[0] >= m2[0]:
                continue
            if m1[0] in groups[0] and m2[0] in groups[0]:
                i1 = proj.image_list[m1[0]]
                i2 = proj.image_list[m2[0]]
                ned1, ypr1, q1 = i1.get_camera_pose(opt=True)
                ned2, ypr2, q2 = i2.get_camera_pose(opt=True)
                quick_approx = False
                if quick_approx:
                    # quick hack angle approximation
                    avg = (np.array(ned1) + np.array(ned2)) * 0.5
                    y = np.linalg.norm(np.array(ned2) - np.array(ned1))
                    x = np.linalg.norm(avg - np.array(match[0]))
                    angle_deg = math.atan2(y, x) * r2d
                else:
                    angle_deg = compute_angle(ned1, ned2, match[0]) * r2d
                pair_angles[m1[0]][m2[0]].append(angle_deg)

print("Computing per-image statistics...")
by_pair = []
for i in range(len(proj.image_list)):
    for j in range(len(proj.image_list)):
        if len(pair_angles[i][j]):
            angles = np.array(pair_angles[i][j])
            avg = np.mean(angles)
            std = np.std(angles)
            min = np.amin(angles)
            #print(i, j, 'avg:', avg, 'std:', std, 'min:', min)
            by_pair.append( [i, j, avg, std, min] )

# (Average angle) pairs with very small average angles between each feature
# and camera location indicate closely located camera poses and these
# cause problems because very small changes in camera pose lead to
# very large changes in feature location.

# (Standard Deviaion) I haven't wrapped my head around this but pairs
# with a high standard deviation of feature to camera pose angles also
# seem to be tied to optimizer/solution degeneracy.

# (Minimum angle) Pairs with a few small minimum angles in the set suggest
# one image is above the other and again run into the problem where
# small pose changes can lead to large feature location changes (but
# just for a few features, not all the features in the pairing.)

avg_cutoff_deg = 2
min_cutoff_deg = 0.5
std_cutoff_deg = 10
print("Marking small angle image pairs for deletion...")
mark_list = []
by_pair = sorted(by_pair, key=lambda fields: fields[4], reverse=False) # by min
for line in by_pair:
    print(line[0], line[1], 'avg: %.2f' % line[2], 'std: %.2f' % line[3], 'min: %.2f' % line[4])
    if line[2] < avg_cutoff_deg or line[3] > std_cutoff_deg or line[4] < min_cutoff_deg:   # cutoff angles (deg)
        print('  (remove)')
        for k, match in enumerate(matches_opt):
            size = len(match[1:])
            for i, m1 in enumerate(match[1:]):
                for j, m2 in enumerate(match[1:]):
                    # this will probably create double markings, but
                    # that's ok, I want to be sure.
                    if m1[0] == line[0] and m2[0] == line[1]:
                        mark_list.append( [k, i] )
                        mark_list.append( [k, j] )
                    if m1[0] == line[1] and m2[0] == line[0]:
                        mark_list.append( [k, i] )
                        mark_list.append( [k, j] )
result = input('Press enter to continue:')

# mark selection
cull.mark_using_list(mark_list, matches_orig)
cull.mark_using_list(mark_list, matches_opt)
mark_sum = len(mark_list)

def delete_marked_features(matches):
    print(" deleting marked items...")
    for i in reversed(range(len(matches))):
        match = matches[i]
        has_bad_elem = False
        for j in reversed(range(1, len(match))):
            p = match[j]
            if p == [-1, -1]:
                has_bad_elem = True
                match.pop(j)
        if len(match) < 3:
            print("deleting match that is now in less than 2 images:", match)
            matches.pop(i)

mark_sum = len(mark_list)
if mark_sum > 0:
    print('Outliers to remove from match lists:', mark_sum)
    result = input('Save these changes? (y/n):')
    if result == 'y' or result == 'Y':
        delete_marked_features(matches_orig)
        delete_marked_features(matches_opt)
        # write out the updated match dictionaries
        print("Writing original matches...")
        pickle.dump(matches_orig, open(os.path.join(args.project, source), "wb"))
        print("Writing optimized matches...")
        pickle.dump(matches_opt, open(os.path.join(args.project, "matches_opt"), "wb"))

max_stars_repo_name: skywalkerisnull/ImageAnalysis

max_issues_repo_name: skywalkerisnull/ImageAnalysis

max_forks_repo_name: skywalkerisnull/ImageAnalysis


text:
import sys
import math
import os

import torch
import torchvision
import numpy as np

from pkg_resources import resource_stream

def interpolate1d(x, values, tangents):
    '''
    Returns:
        Returns the interpolated or extrapolated values for each query point,
        depending on whether or not the query lies within the span of the spline.
    '''
    assert torch.is_tensor(x)
    assert torch.is_tensor(values)
    assert torch.is_tensor(tangents)
    float_dtype = x.dtype
    assert values.dtype == float_dtype
    assert tangents.dtype == float_dtype
    assert len(values.shape) == 1
    assert len(tangents.shape) == 1
    assert values.shape[0] == tangents.shape[0]

    x_lo = torch.floor(torch.clamp(x, torch.as_tensor(0),
                                   values.shape[0] - 2)).type(torch.int64)
    x_hi = x_lo + 1

    # Compute the relative distance between each `x` and the knot below it.
    t = x - x_lo.type(float_dtype)

    # Compute the cubic hermite expansion of `t`.
    t_sq = t**2
    t_cu = t * t_sq
    h01 = -2. * t_cu + 3. * t_sq
    h00 = 1. - h01
    h11 = t_cu - t_sq
    h10 = h11 - t_sq + t

    # Linearly extrapolate above and below the extents of the spline for all
    # values.
    value_before = tangents[0] * t + values[0]
    value_after = tangents[-1] * (t - 1.) + values[-1]

    # Cubically interpolate between the knots below and above each query point.
    neighbor_values_lo = values[x_lo]
    neighbor_values_hi = values[x_hi]
    neighbor_tangents_lo = tangents[x_lo]
    neighbor_tangents_hi = tangents[x_hi]
    value_mid = (
        neighbor_values_lo * h00 + neighbor_values_hi * h01 +
        neighbor_tangents_lo * h10 + neighbor_tangents_hi * h11)

    return torch.where(t < 0., value_before,
                       torch.where(t > 1., value_after, value_mid))

def log_safe(x):
    x = torch.as_tensor(x)
    return torch.log(torch.min(x, torch.tensor(33e37).to(x)))

def load_spline_params():
    dirname = os.path.dirname(__file__)
    with open(os.path.join(dirname, '../misc/partition_spline.npz'), "rb") as spline_file:
        with np.load(spline_file, allow_pickle=False) as f:
            spline_x_scale = torch.tensor(f['x_scale'])
            spline_values = torch.tensor(f['values'])
            spline_tangents = torch.tensor(f['tangents'])

    return spline_x_scale, spline_values, spline_tangents

def get_partition_init(shape):
    shape = torch.as_tensor(shape)

    base1 = (2.25 * shape - 4.5) / (torch.abs(shape - 2) + 0.25) + shape + 2
    base2 = 5. / 18. * log_safe(4 * shape - 15) + 8

    return torch.where(shape < 4, base1, base2)

def get_partition(shape):
    shape = torch.as_tensor(shape)
    assert (shape >= 0).all()

    init = get_partition_init(shape)

    x_scale, values, tangents = load_spline_params()

    return interpolate1d(init * x_scale.to(init), values.to(init), tangents.to(init))

def general_adaptive_loss(x, shape, bowl=1.):
    input_shape = x.shape
    shape = torch.as_tensor(shape).to(x.device)
    bowl = torch.as_tensor(bowl).to(x.device)

    b = x.size(0)
    x = x.view(b, -1)

    if len(shape.shape) == 0:
        shape = shape.unsqueeze(dim=0).expand([b, ]).unsqueeze(dim=1)
    else:
        shape = shape.view(b, -1)

    if len(bowl.shape) == 0:
        bowl = bowl.unsqueeze(dim=0).expand([b, ]).unsqueeze(dim=1)
    else:
        bowl = bowl.view(b, -1)

    partition = get_partition(shape)
    ans = (torch.abs(shape - 2)/shape) * (torch.pow((torch.square(x/bowl) /
                                                     torch.abs(shape - 2) + 1), shape/2) - 1) + log_safe(bowl) + log_safe(partition)

    return ans.view(input_shape)

max_stars_repo_name: jmendozais/SDSSDepth

max_issues_repo_name: jmendozais/SDSSDepth

max_forks_repo_name: jmendozais/SDSSDepth


text:
# Flow visualization code used from https://github.com/tomrunia/OpticalFlow_Visualization

# MIT License
#
# Copyright (c) 2018 Tom Runia
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to conditions.
#
# Author: Tom Runia
# Date Created: 2018-08-03

import numpy as np

def make_colorwheel():
    """
    Generates a color wheel for optical flow visualization as presented in:
        Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
        URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf

    Code follows the original C++ source code of Daniel Scharstein.
    Code follows the the Matlab source code of Deqing Sun.

    Returns:
        np.ndarray: Color wheel
    """

    RY = 15
    YG = 6
    GC = 4
    CB = 11
    BM = 13
    MR = 6

    ncols = RY + YG + GC + CB + BM + MR
    colorwheel = np.zeros((ncols, 3))
    col = 0

    # RY
    colorwheel[0:RY, 0] = 255
    colorwheel[0:RY, 1] = np.floor(255*np.arange(0,RY)/RY)
    col = col+RY
    # YG
    colorwheel[col:col+YG, 0] = 255 - np.floor(255*np.arange(0,YG)/YG)
    colorwheel[col:col+YG, 1] = 255
    col = col+YG
    # GC
    colorwheel[col:col+GC, 1] = 255
    colorwheel[col:col+GC, 2] = np.floor(255*np.arange(0,GC)/GC)
    col = col+GC
    # CB
    colorwheel[col:col+CB, 1] = 255 - np.floor(255*np.arange(CB)/CB)
    colorwheel[col:col+CB, 2] = 255
    col = col+CB
    # BM
    colorwheel[col:col+BM, 2] = 255
    colorwheel[col:col+BM, 0] = np.floor(255*np.arange(0,BM)/BM)
    col = col+BM
    # MR
    colorwheel[col:col+MR, 2] = 255 - np.floor(255*np.arange(MR)/MR)
    colorwheel[col:col+MR, 0] = 255
    return colorwheel

def flow_uv_to_colors(u, v, convert_to_bgr=False):
    """
    Applies the flow color wheel to (possibly clipped) flow components u and v.

    According to the C++ source code of Daniel Scharstein
    According to the Matlab source code of Deqing Sun

    Args:
        u (np.ndarray): Input horizontal flow of shape [H,W]
        v (np.ndarray): Input vertical flow of shape [H,W]
        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.

    Returns:
        np.ndarray: Flow visualization image of shape [H,W,3]
    """
    flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
    colorwheel = make_colorwheel()  # shape [55x3]
    ncols = colorwheel.shape[0]
    rad = np.sqrt(np.square(u) + np.square(v))
    a = np.arctan2(-v, -u)/np.pi
    fk = (a+1) / 2*(ncols-1)
    k0 = np.floor(fk).astype(np.int32)
    k1 = k0 + 1
    k1[k1 == ncols] = 0
    f = fk - k0
    for i in range(colorwheel.shape[1]):
        tmp = colorwheel[:,i]
        col0 = tmp[k0] / 255.0
        col1 = tmp[k1] / 255.0
        col = (1-f)*col0 + f*col1
        idx = (rad <= 1)
        col[idx]  = 1 - rad[idx] * (1-col[idx])
        col[~idx] = col[~idx] * 0.75   # out of range
        # Note the 2-i => BGR instead of RGB
        ch_idx = 2-i if convert_to_bgr else i
        flow_image[:,:,ch_idx] = np.floor(255 * col)
    return flow_image

def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False):
    """
    Expects a two dimensional flow image of shape.

    Args:
        flow_uv (np.ndarray): Flow UV image of shape [H,W,2]
        clip_flow (float, optional): Clip maximum of flow values. Defaults to None.
        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.

    Returns:
        np.ndarray: Flow visualization image of shape [H,W,3]
    """
    assert flow_uv.ndim == 3, 'input flow must have three dimensions'
    assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
    if clip_flow is not None:
        flow_uv = np.clip(flow_uv, 0, clip_flow)
    u = flow_uv[:,:,0]
    v = flow_uv[:,:,1]
    rad = np.sqrt(np.square(u) + np.square(v))
    rad_max = np.max(rad)
    epsilon = 1e-5
    u = u / (rad_max + epsilon)
    v = v / (rad_max + epsilon)
    return flow_uv_to_colors(u, v, convert_to_bgr)

max_stars_repo_name: hologerry/UPFlow

max_issues_repo_name: hologerry/UPFlow

max_forks_repo_name: hologerry/UPFlow


text:
import pandas as pd
from pandas import DataFrame
import numpy as np
import csv
import requests
from bs4 import BeautifulSoup
from tqdm import trange
pd.set_option('display.max_columns', 999)

#this entire program shoudl be rewritten using basketball refrence scraper API
#nba is too slow

def main():
    def_rating = career_def_rating()
    off_rating = career_off_rating()
    def_playoff_rating = career_def_playoffs()
    off_playoff_rating = career_off_playoffs()
    winshare_48 = winshare_per_48()
    winshare_48_playoff = winshare_48_playoffs()

    def_rating.to_csv("/Users/lucasmunson/Desktop/nokogiri tutorial/GOAT Calc/player_data/def_rating.csv")
    off_rating.to_csv("/Users/lucasmunson/Desktop/nokogiri tutorial/GOAT Calc/player_data/off_rating.csv")
    def_playoff_rating.to_csv("/Users/lucasmunson/Desktop/nokogiri tutorial/GOAT Calc/player_data/def_playoff_rating.csv")
    off_playoff_rating.to_csv("/Users/lucasmunson/Desktop/nokogiri tutorial/GOAT Calc/player_data/off_playoff_rating.csv")
    winshare_48.to_csv("/Users/lucasmunson/Desktop/nokogiri tutorial/GOAT Calc/player_data/winshare_48.csv")
    winshare_48_playoff.to_csv("/Users/lucasmunson/Desktop/nokogiri tutorial/GOAT Calc/player_data/winshare_48_playoff.csv")

def career_def_rating():
    page = requests.get('https://www.basketball-reference.com/leaders/def_rtg_career.html')
    df_list = pd.read_html(page.content)

    def_rating = df_list[0]
    def_rating['Player'] = def_rating['Player'].str.replace('*','')
    return def_rating

def career_off_rating():
    page = requests.get('https://www.basketball-reference.com/leaders/off_rtg_career.html')
    off_list = pd.read_html(page.content)

    off_rating = off_list[0]
    off_rating['Player'] = off_rating['Player'].str.replace('*','')
    return off_rating

def career_def_playoffs():
    page = requests.get('https://www.basketball-reference.com/leaders/def_rtg_career_p.html')
    def_playoffs_list = pd.read_html(page.content)

    def_playoff_rating = def_playoffs_list[0]
    def_playoff_rating['Player'] = def_playoff_rating['Player'].str.replace('*','')
    return def_playoff_rating

def career_off_playoffs():
    page = requests.get('https://www.basketball-reference.com/leaders/off_rtg_career_p.html')
    off_playoffs_list = pd.read_html(page.content)

    off_playoff_rating = off_playoffs_list[0]
    off_playoff_rating['Player'] = off_playoff_rating['Player'].str.replace('*','')
    return off_playoff_rating

def winshare_per_48():
    page = requests.get('https://www.basketball-reference.com/leaders/ws_per_48_career.html')
    winshare_48_list = pd.read_html(page.content)

    winshare_48 = winshare_48_list[0]
    winshare_48['Player'] = winshare_48['Player'].str.replace('*','')
    return winshare_48

def winshare_48_playoffs():
    page = requests.get('https://www.basketball-reference.com/leaders/ws_per_48_career_p.html')
    winshare_48_playoff = pd.read_html(page.content)

    winshare_48_playoff = winshare_48_playoff[0]
    winshare_48_playoff['Player'] = winshare_48_playoff['Player'].str.replace('*','')
    return winshare_48_playoff

if __name__ == "__main__":
    main()

max_stars_repo_name: Munseaux/GOAT-calc

max_issues_repo_name: Munseaux/GOAT-calc

max_forks_repo_name: Munseaux/GOAT-calc


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

def plotStimRasters(stimulus, samples, spikes, unit, ltime, rtime, save=False, saveString = '',
                    baseline=0, sample_rate=20000, fig_size=(10,4),
                    heightRatio=[1,4], markerSize=3,stimLabel='Force (mN)'):
    """
    Generate plots with stimulus displayed at top and rasters below for individual units.

    inputs:
        stimulus - list of ndarrays, stimulus waveform
        samples - list of ndarrays, spike times in samples
        spikes - list of ndarrays, cluster identity for each spike
        unit - int, unit to include in raster
        save - boolean, save plot to disk, default=False
        baseline - float, time before first stimulus in s, default=0.0
        sample_rate - int, sample rate in Hz, default=20000
        fig_size - tuple, ratio of width to length
        heightRatio - list, ratio of heights of stimulus and raster plots

    generates a plot; no outputs
    """

    # Plot stimulus waveform
    f, (a0, a1) = plt.subplots(2,1,gridspec_kw={'height_ratios':heightRatio},figsize=fig_size)

    if ltime >= 0:
        xaxis = np.arange(ltime-baseline,rtime-baseline,1/sample_rate)
        for sweep in stimulus:
            a0.plot(xaxis,sweep[int(sample_rate*ltime):int(sample_rate*rtime)],linewidth=.5,color='r') # add +5*i to the y axis to get separate traces

    xlim = [ltime,rtime]
    a0.set_xlim(xlim)
    a0.set_title('Unit '+str(unit))
    a0.set_xticks([])
    a0.set_ylabel(stimLabel)
    a0.spines['top'].set_visible(False)
    a0.spines['right'].set_visible(False)
    a0.spines['bottom'].set_visible(False)

    # Plot Rasters
    for sweep in range(len(samples)):
        sweepspikes = spikes[sweep][spikes[sweep]==unit]
        sweepsamples = samples[sweep][spikes[sweep]==unit]
        sweepspikes = sweepspikes[(sweepsamples > ltime*sample_rate) & (sweepsamples < rtime*sample_rate)]
        sweepsamples = sweepsamples[(sweepsamples > ltime*sample_rate) & (sweepsamples < rtime*sample_rate)]
        a1.plot(sweepsamples/sample_rate-baseline,(sweepspikes+sweep-unit),'|',color='k',markersize=markerSize,mew=.5)
    a1.set_xlim(xlim)
    a1.set_ylim(-1,len(samples))
    a1.set_xlabel('Time (s)')
    a1.set_ylabel('Trial')
    a1.spines['top'].set_visible(False)
    a1.spines['right'].set_visible(False)
    a1.invert_yaxis()
    if save:
        plt.savefig('RasterUnit'+str(unit)+saveString+'.png',dpi=300, bbox_inches='tight')
    plt.show()
    plt.close()

def makeSweepPSTH(bin_size, samples, spikes,sample_rate=20000, units=None, duration=None, verbose=False, rate=True, bs_window=[0, 0.25]):
    """
    Use this to convert spike time rasters into PSTHs with user-defined bin

    inputs:
        bin_size - float, bin size in seconds
        samples - list of ndarrays, time of spikes in samples
        spikes- list of ndarrays, spike cluster identities
        sample_rate - int, Hz, default = 20000
        units - None or sequence, list of units to include in PSTH
        duration - None or float, duration of PSTH; if None, inferred from last spike
        verbose - boolean, print information about psth during calculation
        rate - boolean; Output rate (divide by bin_size and # of trials) or total spikes per trial (divide by # trials only)
        bs_window - sequence, len 2; window (in s) to use for baseline subtraction; default = [0, 0.25]
    output: dict with keys:
        psths - ndarray
        bin_size - float, same as input
        sample_rate - int, same as input
        xaxis - ndarray, gives the left side of the bins
        units - ndarray, units included in psth
    """

    bin_samples = bin_size * sample_rate

    if duration is None:
        maxBin = max(np.concatenate(samples))/sample_rate
        maxBin = np.ceil(maxBin/bin_size)*bin_size
        minTime = min(np.concatenate(samples))/sample_rate
        minBin = np.floor(minTime/bin_size)*bin_size
    else:
        minTime = min(np.concatenate(samples))/sample_rate
        if minTime < 0:
            minBin = np.floor(minTime/bin_size)*bin_size
        else:
            minBin = 0
        maxBin = duration - minBin

    if minBin < 0:
        modSamples = [np.array(n)-minTime*sample_rate for n in samples]
    else:
        modSamples = samples

    if units is None:  # if user does not specify which units to use (usually done with np.unique(goodSpikes))
        units = np.unique(np.hstack(spikes))
    numUnits = len(units)

    psths = np.zeros([int(np.ceil((maxBin-minBin)/bin_size)), numUnits])
    if verbose:
        print('psth size is',psths.shape)
    for i in range(len(modSamples)):
        for stepSample, stepSpike in zip(modSamples[i], spikes[i]):
            if stepSpike in units:
                if int(np.floor(stepSample/bin_samples)) == psths.shape[0]:
                    psths[int(np.floor(stepSample/bin_samples))-1, np.where(units == stepSpike)[0][0]] += 1 ## for the rare instance when a spike is detected at the last sample of a sweep
                else:
                    psths[int(np.floor(stepSample/bin_samples)), np.where(units == stepSpike)[0][0]] += 1
    psth_dict = {}
    if rate:
        psth_dict['psths'] = psths/bin_size/len(modSamples) # in units of Hz
    else:
        psth_dict['psths'] = psths/len(modSamples) # in units of spikes/trial in each bin

    psths_bs = np.copy(np.transpose(psth_dict['psths']))
    psths_z = np.copy(np.transpose(psth_dict['psths']))
    for i,psth in enumerate(psths_bs):
        tempMean = np.mean(psth[int((bs_window[0]-minBin)/bin_size):int((bs_window[1]-minBin)/bin_size)])
        tempStDev = np.std(psth[int((bs_window[0]-minBin)/bin_size):int((bs_window[1]-minBin)/bin_size)])
        #print(tempMean)
        psths_bs[i] = psth - tempMean
        psths_z[i] = psths_bs[i]/tempStDev
    psth_dict['psths_bs'] = np.transpose(psths_bs)
    psth_dict['psths_z'] = np.transpose(psths_z)
    psth_dict['bin_size'] = bin_size # in s
    psth_dict['sample_rate'] = sample_rate # in Hz
    psth_dict['xaxis'] = np.arange(minBin,maxBin,bin_size)
    psth_dict['units'] = units
    psth_dict['num_sweeps'] = len(samples)
    return psth_dict

def singleSweepPSTH(bin_size, samples, spikes,sample_rate=20000, units=None, duration=None, verbose=False, rate=True, bs_window=[0, 0.25]):
    """
    Use this to convert spike time rasters into PSTHs with user-defined bin

    inputs:
        bin_size - float, bin size in seconds
        samples - list of ndarrays, time of spikes in samples
        spikes- list of ndarrays, spike cluster identities
        sample_rate - int, Hz, default = 20000
        units - None or sequence, list of units to include in PSTH
        duration - None or float, duration of PSTH; if None, inferred from last spike
        verbose - boolean, print information about psth during calculation
        rate - boolean; Output rate (divide by bin_size and # of trials) or total spikes per trial (divide by # trials only)
        bs_window - sequence, len 2; window (in s) to use for baseline subtraction; default = [0, 0.25]
    output: dict with keys:
        psths - ndarray
        bin_size - float, same as input
        sample_rate - int, same as input
        xaxis - ndarray, gives the left side of the bins
        units - ndarray, units included in psth
    """

    bin_samples = bin_size * sample_rate

    if duration is None:
        maxBin = max(np.concatenate(samples))/sample_rate
    else:
        maxBin = duration

    if units is None:  # if user does not specify which units to use (usually done with np.unique(goodSpikes))
        units = np.unique(np.hstack(spikes))
    numUnits = len(units)
    numSweeps = len(samples)

    psths = np.zeros([int(np.ceil(maxBin/bin_size)), numSweeps, numUnits])
    if verbose:
        print('psth size is',psths.shape)
    for i in range(numSweeps):
        for stepSample, stepSpike in zip(samples[i], spikes[i]):
            if stepSpike in units:
                if int(np.floor(stepSample/bin_samples)) == psths.shape[0]:
                    psths[int(np.floor(stepSample/bin_samples))-1, i, np.where(units == stepSpike)[0][0]] += 1 ## for the rare instance when a spike is detected at the last sample of a sweep
                else:
                    psths[int(np.floor(stepSample/bin_samples)), i, np.where(units == stepSpike)[0][0]] += 1
    psth_dict = {}
    if rate:
        psth_dict['psths'] = psths/bin_size # in units of Hz
    else:
        psth_dict['psths'] = psths # in units of spikes/trial in each bin

    psths_mean = np.mean(psth_dict['psths'],axis=1) ## collapse to mean across sweeps
    psths_sd = np.std(psth_dict['psths'],axis=1)
    psths_bs = np.copy(np.transpose(psths_mean))
    psths_z = np.copy(np.transpose(psths_mean))
    for i,psth in enumerate(psths_bs):
        tempMean = np.mean(psth[int(bs_window[0]/bin_size):int(bs_window[1]/bin_size)])
        tempStDev = np.std(psth[int(bs_window[0]/bin_size):int(bs_window[1]/bin_size)])
        #print(tempMean)
        psths_bs[i] = psth - tempMean
        psths_z[i] = (psth - tempMean)/tempStDev
    psth_dict['psths_mean'] = psths_mean
    psth_dict['psths_sd'] = psths_sd
    psth_dict['psths_sem'] = psths_sd/np.sqrt(psths.shape[1])
    psth_dict['psths_bs'] = np.transpose(psths_bs)
    psth_dict['psths_z'] = np.transpose(psths_z)
    psth_dict['bin_size'] = bin_size # in s
    psth_dict['sample_rate'] = sample_rate # in Hz
    psth_dict['xaxis'] = np.arange(0,maxBin,bin_size)
    psth_dict['units'] = units
    psth_dict['num_sweeps'] = numSweeps
    return psth_dict

def spikeTriggeredAverage(spikes, stimulusTrace, window=(-.1,.1), sampleRate=20000):
    """
    Calculates the average of a stimulus trace around each spike

    Inputs:
        spikes - sequence - list of spike times to trigger against
        stimulusTrace - ndarray - stimulus to average
        window - sequence, len 2 - period over which to average in s, defaults to (-0.1, 0.1)
        sampleRate - int - sample rate
    Outputs:
        sta - sta from -window to window at rate of sampleRate
        xaxis - ndarray, sequence for xaxis of sta
    """
    sta = np.zeros(int((window[1]-window[0])*sampleRate))
    window_samples = [int(n*sampleRate) for n in window]
    numErrors = 0
    for spike in spikes:
        try:
            sta = sta + stimulusTrace[int(spike+window_samples[0]):int(spike+window_samples[1])]
        except ValueError:
            numErrors += 1
    sta /= len(spikes) - numErrors
    xaxis = np.arange(window[0],window[1],1/sampleRate)

    return sta, xaxis

max_stars_repo_name: ajemanuel/analyzeMEA

max_issues_repo_name: ajemanuel/analyzeMEA

max_forks_repo_name: ajemanuel/analyzeMEA


text:
import os
import glob
import numpy as np

v_max = 5
v_min = 0.5
points = 15
i = 0
root = os.path.abspath(os.path.dirname(__file__))
os.system('rm -rf %s/raw/group_velocity/*'%root)
os.system('rm -rf %s/raw/phase_velocity/*'%root)
os.system('rm -rf %s/pick_result/group_velocity/*'%root)
os.system('rm -rf %s/pick_result/phase_velocity/*'%root)
for file in os.listdir(root + '/test_result'):
    i += 1
    print(i, file)
    path = os.path.join(root + '/test_result', file)
    os.chdir(path)
    for group in glob.glob('newG*.txt'):
        name = group[5:]
        os.system('cp -rf ./%s %s/raw/group_velocity/%s'%(group, root, name))
    for phase in glob.glob('newC*.txt'):
        name = phase[5:]
        os.system('cp -rf ./%s %s/raw/phase_velocity/%s'%(phase, root, name))

os.chdir(root + '/raw/group_velocity')
all_ = 0
for file in glob.glob('*txt'):
    all_ += 1
    print(all_, file)
    num = 0
    data = np.loadtxt(file)[:, 1]
    for i in data:
        if i <= v_max and i >= v_min:
            num += 1
    if num >= points:
        # print(num)
        os.system('cp -rf %s %s/pick_result/group_velocity/'%(file, root))
        data2 = np.loadtxt(root + '/raw/phase_velocity/'+file)[:, 1]
        num2 = 0
        for i in data2:
            if i <= v_max and i >= v_min:
                num2 += 1
        if num2 >= points:
            # print(num2)
            os.system('cp -rf %s/raw/phase_velocity/%s %s/pick_result/phase_velocity/'%(root, file, root))

max_stars_repo_name: ShaoboYang-USTC/DisperPicker

max_issues_repo_name: ShaoboYang-USTC/DisperPicker

max_forks_repo_name: ShaoboYang-USTC/DisperPicker