text:
import copy
from abc import abstractmethod

import numpy as np
import torch

from .hooks import HookPhase

IMPLEMENTED_OPTIM = ["adam", "rmsprop", "sgd"]

class BaseTrainer:
    def __init__(
        self,
        optimizer_name="adam",
        lr=1e-4,
        n_epochs=100,
        betas=[0.9, 0.99],
        batch_size=512,
        log_every=5,
        checkpoint=20,
        device="cuda",
        hooks=None,
        momentum=0,
        rms_alpha=0.99,
        target_mean=None,
        target_std=None,
        early_stop_metric=None,
        early_stop_iters=None,
        track_best_metrics=None,
        log_intermediate_results=False,
    ):
        if optimizer_name not in IMPLEMENTED_OPTIM:
            raise ValueError(f"Unkown optmizer name: {optimizer_name}.")

        self.optimizer_name = optimizer_name
        self.model_regularizer = None  # Lazy init
        self.lr = lr
        self.n_epochs = n_epochs
        self.momentum = momentum
        self.batch_size = batch_size
        self.betas = (betas[0], betas[1])
        self.rms_alpha = rms_alpha

        if hooks is not None:
            self.hooks = hooks
        else:
            self.hooks = []

        self.checkpoint = checkpoint
        self.log_every = log_every
        self.log_intermediate_results = log_intermediate_results

        if device == "cuda":
            if not torch.cuda.is_available():
                raise ValueError("CUDA requested, but it's not available")
            self.device = device
        else:
            self.device = "cpu"

        self._target_mean = None
        self._target_std = None

        self.early_stop_metric = early_stop_metric
        self.early_stop_iters = early_stop_iters
        self.early_stop_model = None
        self.early_stop_best = np.inf
        self._early_stop_no_improvements = 0
        self.track_best_metrics = track_best_metrics if track_best_metrics else {}
        self.track_best_models = {}
        self._track_best_no_improvements = {m: 0 for m in self.track_best_metrics}
        self._track_best_values = {m: np.inf for m in self.track_best_metrics}

    def set_target_normalizer(self, mean, std):
        self._target_mean = torch.from_numpy(np.array(mean)).float().to(self.device)
        self._target_std = torch.from_numpy(np.array(std)).float().to(self.device)

    def _get_optimizer(self, params):
        if self.optimizer_name == "adam":
            optimizer = torch.optim.Adam(params=params, lr=self.lr)
        elif self.optimizer_name == "rmsprop":
            optimizer = torch.optim.RMSprop(
                params=params, lr=self.lr, momentum=self.momentum, alpha=self.rms_alpha
            )
        elif self.optimizer_name == "sgd":
            optimizer = torch.optim.SGD(
                params=params, lr=self.lr, momentum=self.momentum
            )
        else:
            raise ValueError(f"Unkown optimizer name {self.optimizer_name}.")

        return optimizer

    def val_dataset_to_torch(self, val_dataset):
        val_data = val_dataset[:]
        if len(val_data) == 3:
            eval_x, eval_y, _ = val_data
        else:
            eval_x, eval_y = val_data

        if not isinstance(eval_x, torch.Tensor):
            eval_x = torch.from_numpy(eval_x).float().to(self.device)
        if not isinstance(eval_y, torch.Tensor):
            eval_y = torch.from_numpy(eval_y).float().to(self.device)

        return eval_x, eval_y

    def denorm(self, mean, var=None):
        if self._target_mean is None:
            return mean if var is None else (mean, var)
        elif var is not None:
            mean = mean * self._target_std + self._target_mean
            var = var * (self._target_std ** 2)
            return mean, var
        else:
            return mean * self._target_std + self._target_mean

    @torch.no_grad()
    def eval_mse(self, model, inp, target, pred=None):
        inp = inp.to(self.device)
        target = target.to(self.device)

        if pred is None:
            if model.training:
                model.eval()
            pred = model(inp)

        if isinstance(pred, tuple):
            pred = pred[0]

        pred = self.denorm(pred)
        mse = (pred - target) ** 2

        return mse.mean()

    def early_stopping(self, model, logger, epoch, **kwargs):
        def update(no_improvements, best_value, current_value):
            if current_value < best_value:
                return 0, current_value
            else:
                return no_improvements + 1, best_value

        for metric in self.track_best_metrics:
            if metric not in kwargs:
                logger.info(
                    f"Warning: requested tracking best metric {metric}, but was not found"
                )
                continue
            no_improvements, best_value = update(
                self._track_best_no_improvements[metric],
                self._track_best_values[metric],
                kwargs[metric],
            )
            self._track_best_no_improvements[metric] = no_improvements
            self._track_best_values[metric] = best_value
            if no_improvements == 0:
                print(f"Saving best {metric}: {best_value}")
                self.track_best_models[metric] = copy.deepcopy(model)
                self.save_checkpoint(
                    logger, model, epoch, postfix=f"best_{metric}", save_hooks=False
                )

        if self.early_stop_metric is None:
            return False

        if self.early_stop_metric not in kwargs:
            raise ValueError(
                "Did not find early stopping metric {self.early_stop_metric}"
            )

        current_value = kwargs[self.early_stop_metric]

        self._early_stop_no_improvements, self.early_stop_best = update(
            self._early_stop_no_improvements, self.early_stop_best, current_value
        )
        if self._early_stop_no_improvements == 0:
            self.early_stop_model = copy.deepcopy(model)
            self.save_checkpoint(logger, model, epoch, postfix="best", save_hooks=False)

        should_early_stop = self._early_stop_no_improvements >= self.early_stop_iters

        if should_early_stop:
            early_stop_epoch = epoch - self._early_stop_no_improvements
            self.execute_hooks_for_phase(
                HookPhase.EARLY_STOP,
                self.early_stop_model,
                logger,
                early_stop_epoch,
                self.early_stop_metric,
                self.early_stop_best,
            )

        return should_early_stop

    @abstractmethod
    def train():
        raise NotImplementedError("Method train not implemented.")

    @abstractmethod
    def eval_model(self, model, xs):
        raise NotImplementedError("Method eval_model not implemented.")

    def init_model(self, model):
        model.to(self.device)

        if hasattr(model, "regularizer"):
            self.model_regularizer = model.regularizer

        return model

    def backward(self, step, loss, mean_loss, optimizer):
        mean_loss.backward()

    def _gather_gradients(self, model):
        gradients = []
        for p in model.parameters():
            if p.grad is None:
                gradients.append(torch.zeros_like(p).view(-1))
            else:
                gradients.append(p.grad.view(-1))

        return torch.cat(gradients)

    def _eval_gradient(self, model, optimizer, loss, with_update=False):
        optimizer.zero_grad()
        loss.backward(retain_graph=not with_update)

        gradients = self._gather_gradients(model)

        if with_update:
            with torch.no_grad():
                params_before = [p.cpu().clone() for p in model.parameters()]
                optimizer_state = optimizer.state_dict()

                optimizer.step()

                updates = [
                    (p2 - p1).view(-1)
                    for p1, p2 in zip(params_before, model.parameters())
                ]

                optimizer.load_state_dict(optimizer_state)
                for p_after, p_before in zip(model.parameters(), params_before):
                    p_after.copy_(p_before)

            return gradients, torch.cat(updates)
        else:
            return gradients

    def execute_hooks_for_phase(self, phase: HookPhase, *args, **kwargs):
        for hook in self.hooks:
            if phase in hook.phases:
                hook(phase, *args, **kwargs)

    def log(self, logger, epoch, info_dict=None):
        if info_dict is not None and self.log_intermediate_results:
            logger.log_training_as_dataframe(info_dict)

        self.execute_hooks_for_phase(HookPhase.LOGGING, logger, epoch)

    def save_checkpoint(self, logger, model, epoch, postfix=None, save_hooks=True):
        postfix = str(epoch + 1) if postfix is None else postfix
        logger.save_model(model, f"m_{postfix}")
        if save_hooks and len(self.hooks) > 0:
            logger.save_hooks(self.hooks)

max_stars_repo_name: martius-lab/beta-nll

max_issues_repo_name: martius-lab/beta-nll

max_forks_repo_name: martius-lab/beta-nll


text:
# -*- coding: utf-8 -*-
from __future__ import division
import numpy
import scipy.stats
import warnings
import numbers
from .hist_tools import SparseAxis, DenseAxis, overflow_behavior, Interval

# Plotting is always terrible
# Let's try our best to follow matplotlib idioms
# https://matplotlib.org/tutorials/introductory/usage.html#coding-styles

_coverage1sd = scipy.stats.norm.cdf(1) - scipy.stats.norm.cdf(-1)

def poisson_interval(sumw, sumw2, coverage=_coverage1sd):
    """Frequentist coverage interval for Poisson-distributed observations

    Parameters
    ----------
        sumw : numpy.ndarray
            Sum of weights vector
        sumw2 : numpy.ndarray
            Sum weights squared vector
        coverage : float, optional
            Central coverage interval, defaults to 68%

    Calculates the so-called 'Garwood' interval,
    c.f. https://www.ine.pt/revstat/pdf/rs120203.pdf or
    http://ms.mcmaster.ca/peter/s743/poissonalpha.html
    For weighted data, this approximates the observed count by ``sumw**2/sumw2``, which
    effectively scales the unweighted poisson interval by the average weight.
    This may not be the optimal solution: see https://arxiv.org/pdf/1309.1287.pdf for a proper treatment.
    When a bin is zero, the scale of the nearest nonzero bin is substituted to scale the nominal upper bound.
    If all bins zero, a warning is generated and interval is set to ``sumw``.
    """
    scale = numpy.empty_like(sumw)
    scale[sumw != 0] = sumw2[sumw != 0] / sumw[sumw != 0]
    if numpy.sum(sumw == 0) > 0:
        missing = numpy.where(sumw == 0)
        available = numpy.nonzero(sumw)
        if len(available[0]) == 0:
            warnings.warn(
                "All sumw are zero!  Cannot compute meaningful error bars",
                RuntimeWarning,
            )
            return numpy.vstack([sumw, sumw])
        nearest = sum(
            [numpy.subtract.outer(d, d0) ** 2 for d, d0 in zip(available, missing)]
        ).argmin(axis=0)
        argnearest = tuple(dim[nearest] for dim in available)
        scale[missing] = scale[argnearest]
    counts = sumw / scale
    lo = scale * scipy.stats.chi2.ppf((1 - coverage) / 2, 2 * counts) / 2.0
    hi = scale * scipy.stats.chi2.ppf((1 + coverage) / 2, 2 * (counts + 1)) / 2.0
    interval = numpy.array([lo, hi])
    interval[interval == numpy.nan] = 0.0  # chi2.ppf produces nan for counts=0
    return interval

def clopper_pearson_interval(num, denom, coverage=_coverage1sd):
    """Compute Clopper-Pearson coverage interval for a binomial distribution

    Parameters
    ----------
        num : numpy.ndarray
            Numerator, or number of successes, vectorized
        denom : numpy.ndarray
            Denominator or number of trials, vectorized
        coverage : float, optional
            Central coverage interval, defaults to 68%

    c.f. http://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval
    """
    if numpy.any(num > denom):
        raise ValueError(
            "Found numerator larger than denominator while calculating binomial uncertainty"
        )
    lo = scipy.stats.beta.ppf((1 - coverage) / 2, num, denom - num + 1)
    hi = scipy.stats.beta.ppf((1 + coverage) / 2, num + 1, denom - num)
    interval = numpy.array([lo, hi])
    interval[:, num == 0.0] = 0.0
    interval[1, num == denom] = 1.0
    return interval

def normal_interval(pw, tw, pw2, tw2, coverage=_coverage1sd):
    """Compute errors based on the expansion of pass/(pass + fail), possibly weighted

    Parameters
    ----------
    pw : numpy.ndarray
        Numerator, or number of (weighted) successes, vectorized
    tw : numpy.ndarray
        Denominator or number of (weighted) trials, vectorized
    pw2 : numpy.ndarray
        Numerator sum of weights squared, vectorized
    tw2 : numpy.ndarray
        Denominator sum of weights squared, vectorized
    coverage : float, optional
        Central coverage interval, defaults to 68%

    c.f. https://root.cern.ch/doc/master/TEfficiency_8cxx_source.html#l02515
    """

    eff = pw / tw

    variance = (pw2 * (1 - 2 * eff) + tw2 * eff ** 2) / (tw ** 2)
    sigma = numpy.sqrt(variance)

    prob = 0.5 * (1 - coverage)
    delta = numpy.zeros_like(sigma)
    delta[sigma != 0] = scipy.stats.norm.ppf(prob, scale=sigma[sigma != 0])

    lo = eff - numpy.minimum(eff + delta, numpy.ones_like(eff))
    hi = numpy.maximum(eff - delta, numpy.zeros_like(eff)) - eff

    return numpy.array([lo, hi])

def plot1d(
    hist,
    ax=None,
    clear=True,
    overlay=None,
    stack=False,
    overflow="none",
    line_opts=None,
    fill_opts=None,
    error_opts=None,
    legend_opts={},
    overlay_overflow="none",
    density=False,
    binwnorm=None,
    order=None,
):
    """Create a 1D plot from a 1D or 2D `Hist` object

    Parameters
    ----------
        hist : Hist
            Histogram with maximum of two dimensions
        ax : matplotlib.axes.Axes, optional
            Axes object (if None, one is created)
        clear : bool, optional
            Whether to clear Axes before drawing (if passed); if False, this function will skip drawing the legend
        overlay : str, optional
            In the case that ``hist`` is 2D, specify the axis of hist to overlay (remaining axis will be x axis)
        stack : bool, optional
            Whether to stack or overlay non-axis dimension (if it exists)
        order : list, optional
            How to order when stacking. Take a list of identifiers.
        overflow : str, optional
            If overflow behavior is not 'none', extra bins will be drawn on either end of the nominal
            axis range, to represent the contents of the overflow bins.  See `Hist.sum` documentation
            for a description of the options.
        line_opts : dict, optional
            A dictionary of options to pass to the matplotlib
            `ax.step <https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.step.html>`_ call
            internal to this function.  Leave blank for defaults.
        fill_opts : dict, optional
            A dictionary of options to pass to the matplotlib
            `ax.fill_between <https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.axes.Axes.fill_between.html>`_ call
            internal to this function.  Leave blank for defaults.
        error_opts : dict, optional
            A dictionary of options to pass to the matplotlib
            `ax.errorbar <https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.errorbar.html>`_ call
            internal to this function.  Leave blank for defaults.  Some special options are interpreted by
            this function and not passed to matplotlib: 'emarker' (default: '') specifies the marker type
            to place at cap of the errorbar.
        legend_opts : dict, optional
            A dictionary of options  to pass to the matplotlib
            `ax.legend <https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.legend.html>`_ call
            internal to this fuction.  Leave blank for defaults.
        overlay_overflow : str, optional
            If overflow behavior is not 'none', extra bins in the overlay axis will be overlayed or stacked,
            to represent the contents of the overflow bins.  See `Hist.sum` documentation for a description of the options.
        density : bool, optional
            If true, convert sum weights to probability density (i.e. integrates to 1 over domain of axis)
            (Note: this option conflicts with ``binwnorm``)
        binwnorm : float, optional
            If true, convert sum weights to bin-width-normalized, with unit equal to supplied value (usually you want to specify 1.)

    Returns
    -------
        ax : matplotlib.axes.Axes
            A matplotlib `Axes <https://matplotlib.org/3.1.1/api/axes_api.html>`_ object
    """
    import mplhep as hep
    import matplotlib.pyplot as plt

    if ax is None:
        ax = plt.gca()
    else:
        if not isinstance(ax, plt.Axes):
            raise ValueError("ax must be a matplotlib Axes object")
        if clear:
            ax.clear()
    if hist.dim() > 2:
        raise ValueError(
            "plot1d() can only support up to two dimensions (one for axis, one to stack or overlay)"
        )
    if overlay is None and hist.sparse_dim() == 1 and hist.dense_dim() == 1:
        overlay = hist.sparse_axes()[0].name
    elif overlay is None and hist.dim() > 1:
        raise ValueError(
            "plot1d() can only support one dimension without an overlay axis chosen"
        )
    if density and binwnorm is not None:
        raise ValueError("Cannot use density and binwnorm at the same time!")
    if binwnorm is not None:
        if not isinstance(binwnorm, numbers.Number):
            raise ValueError(
                "Bin width normalization not a number, but a %r" % binwnorm.__class__
            )
    if line_opts is None and fill_opts is None and error_opts is None:
        if stack:
            fill_opts = {}
        else:
            line_opts = {}
            error_opts = {}

    axis = hist.axes()[0]
    if overlay is not None:
        overlay = hist.axis(overlay)
        if axis == overlay:
            axis = hist.axes()[1]
    if isinstance(axis, SparseAxis):
        raise NotImplementedError("Plot a sparse axis (e.g. bar chart)")
    elif isinstance(axis, DenseAxis):
        ax.set_xlabel(axis.label)
        ax.set_ylabel(hist.label)
        edges = axis.edges(overflow=overflow)
        if order is None:
            identifiers = (
                hist.identifiers(overlay, overflow=overlay_overflow)
                if overlay is not None
                else [None]
            )
        else:
            identifiers = order
        plot_info = {
            "identifier": identifiers,
            "label": list(map(str, identifiers)),
            "sumw": [],
            "sumw2": [],
        }
        for i, identifier in enumerate(identifiers):
            if identifier is None:
                sumw, sumw2 = hist.values(sumw2=True, overflow=overflow)[()]
            elif isinstance(overlay, SparseAxis):
                sumw, sumw2 = hist.integrate(overlay, identifier).values(
                    sumw2=True, overflow=overflow
                )[()]
            else:
                sumw, sumw2 = hist.values(sumw2=True, overflow="allnan")[()]
                the_slice = (
                    i if overflow_behavior(overlay_overflow).start is None else i + 1,
                    overflow_behavior(overflow),
                )
                if hist._idense(overlay) == 1:
                    the_slice = (the_slice[1], the_slice[0])
                sumw = sumw[the_slice]
                sumw2 = sumw2[the_slice]
            plot_info["sumw"].append(sumw)
            plot_info["sumw2"].append(sumw2)

        def w2err(sumw, sumw2):
            err = []
            for a, b in zip(sumw, sumw2):
                err.append(numpy.abs(poisson_interval(a, b) - a))
            return err

        kwargs = None
        if line_opts is not None and error_opts is None:
            _error = None
        else:
            _error = w2err(plot_info["sumw"], plot_info["sumw2"])
        if fill_opts is not None:
            histtype = "fill"
            kwargs = fill_opts
        elif error_opts is not None and line_opts is None:
            histtype = "errorbar"
            kwargs = error_opts
        else:
            histtype = "step"
            kwargs = line_opts
        if kwargs is None:
            kwargs = {}

        hep.histplot(
            plot_info["sumw"],
            edges,
            label=plot_info["label"],
            yerr=_error,
            histtype=histtype,
            ax=ax,
            density=density,
            binwnorm=binwnorm,
            stack=stack,
            **kwargs
        )

        if stack and error_opts is not None:
            stack_sumw = numpy.sum(plot_info["sumw"], axis=0)
            stack_sumw2 = numpy.sum(plot_info["sumw2"], axis=0)
            err = poisson_interval(stack_sumw, stack_sumw2)
            if binwnorm is not None:
                err *= binwnorm / numpy.diff(edges)[None, :]
            opts = {
                "step": "post",
                "label": "Sum unc.",
                "hatch": "///",
                "facecolor": "none",
                "edgecolor": (0, 0, 0, 0.5),
                "linewidth": 0,
            }
            opts.update(error_opts)
            ax.fill_between(
                x=edges,
                y1=numpy.r_[err[0, :], err[0, -1]],
                y2=numpy.r_[err[1, :], err[1, -1]],
                **opts
            )

        if legend_opts is not None:
            _label = overlay.label if overlay is not None else ""
            ax.legend(title=_label, **legend_opts)
        else:
            ax.legend(title=_label)
        ax.autoscale(axis="x", tight=True)
        ax.set_ylim(0, None)

    return ax

def plotratio(
    num,
    denom,
    ax=None,
    clear=True,
    overflow="none",
    error_opts=None,
    denom_fill_opts=None,
    guide_opts=None,
    unc="clopper-pearson",
    label=None,
):
    """Create a ratio plot, dividing two compatible histograms

    Parameters
    ----------
        num : Hist
            Numerator, a single-axis histogram
        denom : Hist
            Denominator, a single-axis histogram
        ax : matplotlib.axes.Axes, optional
            Axes object (if None, one is created)
        clear : bool, optional
            Whether to clear Axes before drawing (if passed); if False, this function will skip drawing the legend
        overflow : str, optional
            If overflow behavior is not 'none', extra bins will be drawn on either end of the nominal
            axis range, to represent the contents of the overflow bins.  See `Hist.sum` documentation
            for a description of the options.
        error_opts : dict, optional
            A dictionary of options to pass to the matplotlib
            `ax.errorbar <https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.errorbar.html>`_ call
            internal to this function.  Leave blank for defaults.  Some special options are interpreted by
            this function and not passed to matplotlib: 'emarker' (default: '') specifies the marker type
            to place at cap of the errorbar.
        denom_fill_opts : dict, optional
            A dictionary of options to pass to the matplotlib
            `ax.fill_between <https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.axes.Axes.fill_between.html>`_ call
            internal to this function, filling the denominator uncertainty band.  Leave blank for defaults.
        guide_opts : dict, optional
            A dictionary of options to pass to the matplotlib
            `ax.axhline <https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.axes.Axes.axhline.html>`_ call
            internal to this function, to plot a horizontal guide line at ratio of 1.  Leave blank for defaults.
        unc : str, optional
            Uncertainty calculation option: 'clopper-pearson' interval for efficiencies; 'poisson-ratio' interval
            for ratio of poisson distributions; 'num' poisson interval of numerator scaled by denominator value
            (common for data/mc, for better or worse).
        label : str, optional
            Associate a label to this entry (note: y axis label set by ``num.label``)

    Returns
    -------
        ax : matplotlib.axes.Axes
            A matplotlib `Axes <https://matplotlib.org/3.1.1/api/axes_api.html>`_ object
    """
    import matplotlib.pyplot as plt

    if ax is None:
        fig, ax = plt.subplots(1, 1)
    else:
        if not isinstance(ax, plt.Axes):
            raise ValueError("ax must be a matplotlib Axes object")
        if clear:
            ax.clear()
    if not num.compatible(denom):
        raise ValueError(
            "numerator and denominator histograms have incompatible axis definitions"
        )
    if num.dim() > 1:
        raise ValueError("plotratio() can only support one-dimensional histograms")
    if error_opts is None and denom_fill_opts is None and guide_opts is None:
        error_opts = {}
        denom_fill_opts = {}

    axis = num.axes()[0]
    if isinstance(axis, SparseAxis):
        raise NotImplementedError("Ratio for sparse axes (labeled axis with errorbars)")
    elif isinstance(axis, DenseAxis):
        ax.set_xlabel(axis.label)
        ax.set_ylabel(num.label)
        edges = axis.edges(overflow=overflow)
        centers = axis.centers(overflow=overflow)

        sumw_num, sumw2_num = num.values(sumw2=True, overflow=overflow)[()]
        sumw_denom, sumw2_denom = denom.values(sumw2=True, overflow=overflow)[()]

        rsumw = sumw_num / sumw_denom
        if unc == "clopper-pearson":
            rsumw_err = numpy.abs(
                clopper_pearson_interval(sumw_num, sumw_denom) - rsumw
            )
        elif unc == "poisson-ratio":
            # poisson ratio n/m is equivalent to binomial n/(n+m)
            rsumw_err = numpy.abs(
                clopper_pearson_interval(sumw_num, sumw_num + sumw_denom) - rsumw
            )
        elif unc == "num":
            rsumw_err = numpy.abs(
                poisson_interval(rsumw, sumw2_num / sumw_denom ** 2) - rsumw
            )
        elif unc == "normal":
            rsumw_err = numpy.abs(
                normal_interval(sumw_num, sumw_denom, sumw2_num, sumw2_denom)
            )
        else:
            raise ValueError("Unrecognized uncertainty option: %r" % unc)

        if error_opts is not None:
            opts = {"label": label, "linestyle": "none"}
            opts.update(error_opts)
            emarker = opts.pop("emarker", "")
            errbar = ax.errorbar(x=centers, y=rsumw, yerr=rsumw_err, **opts)
            plt.setp(errbar[1], "marker", emarker)
        if denom_fill_opts is not None:
            unity = numpy.ones_like(sumw_denom)
            denom_unc = poisson_interval(unity, sumw2_denom / sumw_denom ** 2)
            opts = {"step": "post", "facecolor": (0, 0, 0, 0.3), "linewidth": 0}
            opts.update(denom_fill_opts)
            ax.fill_between(
                edges,
                numpy.r_[denom_unc[0], denom_unc[0, -1]],
                numpy.r_[denom_unc[1], denom_unc[1, -1]],
                **opts
            )
        if guide_opts is not None:
            opts = {"linestyle": "--", "color": (0, 0, 0, 0.5), "linewidth": 1}
            opts.update(guide_opts)
            ax.axhline(1.0, **opts)

    if clear:
        ax.autoscale(axis="x", tight=True)
        ax.set_ylim(0, None)

    return ax

def plot2d(
    hist,
    xaxis,
    ax=None,
    clear=True,
    xoverflow="none",
    yoverflow="none",
    patch_opts=None,
    text_opts=None,
    density=False,
    binwnorm=None,
):
    """Create a 2D plot from a 2D `Hist` object

    Parameters
    ----------
        hist : Hist
            Histogram with two dimensions
        xaxis : str or Axis
            Which of the two dimensions to use as an x axis
        ax : matplotlib.axes.Axes, optional
            Axes object (if None, one is created)
        clear : bool, optional
            Whether to clear Axes before drawing (if passed); if False, this function will skip drawing the legend
        xoverflow : str, optional
            If overflow behavior is not 'none', extra bins will be drawn on either end of the nominal x
            axis range, to represent the contents of the overflow bins.  See `Hist.sum` documentation
            for a description of the options.
        yoverflow : str, optional
            Similar to ``xoverflow``
        patch_opts : dict, optional
            Options passed to the matplotlib `pcolormesh <https://matplotlib.org/api/_as_gen/matplotlib.pyplot.pcolormesh.html>`_
            call internal to this function, to plot a rectangular grid of patches colored according to the bin values.
            Leave empty for defaults.
        text_opts : dict, optional
            Options passed to the matplotlib `text <https://matplotlib.org/api/text_api.html#matplotlib.text.Text>`_
            call internal to this function, to place a text label at each bin center with the bin value.  Special
            options interpreted by this function and not passed to matplotlib: 'format': printf-style float format
            , default '%.2g'.
        density : bool, optional
            If true, convert sum weights to probability density (i.e. integrates to 1 over domain of axis)
            (Note: this option conflicts with ``binwnorm``)
        binwnorm : float, optional
            If true, convert sum weights to bin-width-normalized, with unit equal to supplied value (usually you want to specify 1.)

    Returns
    -------
        ax : matplotlib.axes.Axes
            A matplotlib `Axes <https://matplotlib.org/3.1.1/api/axes_api.html>`_ object
    """
    import matplotlib.pyplot as plt

    if ax is None:
        fig, ax = plt.subplots(1, 1)
    else:
        if not isinstance(ax, plt.Axes):
            raise ValueError("ax must be a matplotlib Axes object")
        if clear:
            ax.clear()
        fig = ax.figure
    if hist.dim() != 2:
        raise ValueError("plot2d() can only support exactly two dimensions")
    if density and binwnorm is not None:
        raise ValueError("Cannot use density and binwnorm at the same time!")
    if binwnorm is not None:
        if not isinstance(binwnorm, numbers.Number):
            raise ValueError(
                "Bin width normalization not a number, but a %r" % binwnorm.__class__
            )
    if patch_opts is None and text_opts is None:
        patch_opts = {}

    xaxis = hist.axis(xaxis)
    yaxis = hist.axes()[1]
    transpose = False
    if yaxis == xaxis:
        yaxis = hist.axes()[0]
        transpose = True
    if isinstance(xaxis, SparseAxis) or isinstance(yaxis, SparseAxis):
        raise NotImplementedError("Plot a sparse axis (e.g. bar chart or labeled bins)")
    else:
        xedges = xaxis.edges(overflow=xoverflow)
        yedges = yaxis.edges(overflow=yoverflow)
        sumw, sumw2 = hist.values(sumw2=True, overflow="allnan")[()]
        if transpose:
            sumw = sumw.T
            sumw2 = sumw2.T
        # no support for different overflow behavior per axis, do it ourselves
        sumw = sumw[overflow_behavior(xoverflow), overflow_behavior(yoverflow)]
        sumw2 = sumw2[overflow_behavior(xoverflow), overflow_behavior(yoverflow)]
        if (density or binwnorm is not None) and numpy.sum(sumw) > 0:
            overallnorm = numpy.sum(sumw) * binwnorm if binwnorm is not None else 1.0
            areas = numpy.multiply.outer(numpy.diff(xedges), numpy.diff(yedges))
            binnorms = overallnorm / (areas * numpy.sum(sumw))
            sumw = sumw * binnorms
            sumw2 = sumw2 * binnorms ** 2

        if patch_opts is not None:
            opts = {"cmap": "viridis"}
            opts.update(patch_opts)
            pc = ax.pcolormesh(xedges, yedges, sumw.T, **opts)
            ax.add_collection(pc)
            if clear:
                fig.colorbar(pc, ax=ax, label=hist.label)
        if text_opts is not None:
            for ix, xcenter in enumerate(xaxis.centers()):
                for iy, ycenter in enumerate(yaxis.centers()):
                    opts = {
                        "horizontalalignment": "center",
                        "verticalalignment": "center",
                    }
                    if patch_opts is not None:
                        opts["color"] = (
                            "black" if pc.norm(sumw[ix, iy]) > 0.5 else "lightgrey"
                        )
                    opts.update(text_opts)
                    txtformat = opts.pop("format", r"%.2g")
                    ax.text(xcenter, ycenter, txtformat % sumw[ix, iy], **opts)

    if clear:
        ax.set_xlabel(xaxis.label)
        ax.set_ylabel(yaxis.label)
        ax.set_xlim(xedges[0], xedges[-1])
        ax.set_ylim(yedges[0], yedges[-1])

    return ax

def plotgrid(
    h,
    figure=None,
    row=None,
    col=None,
    overlay=None,
    row_overflow="none",
    col_overflow="none",
    **plot_opts
):
    """Create a grid of plots, enumerating identifiers on up to 3 axes

    Parameters
    ----------
        h : Hist
            A histogram with up to 3 axes
        figure : matplotlib.figure.Figure, optional
            If omitted, a new figure is created. Otherwise, the axes will be redrawn on this existing figure.
        row : str
            Name of row axis
        col : str
            Name of column axis
        overlay : str
            name of overlay axis
        row_overflow : str, optional
            If overflow behavior is not 'none', extra bins will be drawn on either end of the nominal x
            axis range, to represent the contents of the overflow bins.  See `Hist.sum` documentation
            for a description of the options.
        col_overflow : str, optional
            Similar to ``row_overflow``
        ``**plot_opts`` : kwargs
            The remaining axis of the histogram, after removing any of ``row,col,overlay`` specified,
            will be the plot axis, with ``plot_opts`` passed to the `plot1d` call.

    Returns
    -------
        axes : numpy.ndarray
            An array of matplotlib `Axes <https://matplotlib.org/3.1.1/api/axes_api.html>`_ objects
    """
    import matplotlib.pyplot as plt

    haxes = set(ax.name for ax in h.axes())
    nrow, ncol = 1, 1
    if row:
        row_identifiers = h.identifiers(row, overflow=row_overflow)
        nrow = len(row_identifiers)
        haxes.remove(row)
    if col:
        col_identifiers = h.identifiers(col, overflow=col_overflow)
        ncol = len(col_identifiers)
        haxes.remove(col)
    if overlay:
        haxes.remove(overlay)
    if len(haxes) > 1:
        raise ValueError(
            "More than one dimension left: %s" % (",".join(ax for ax in haxes),)
        )
    elif len(haxes) == 0:
        raise ValueError("Not enough dimensions available in %r" % h)

    figsize = plt.rcParams["figure.figsize"]
    figsize = figsize[0] * max(ncol, 1), figsize[1] * max(nrow, 1)
    if figure is None:
        fig, axes = plt.subplots(
            nrow, ncol, figsize=figsize, squeeze=False, sharex=True, sharey=True
        )
    else:
        fig = figure
        shape = (0, 0)
        lastax = fig.get_children()[-1]
        if isinstance(lastax, plt.Axes):
            shape = lastax.rowNum + 1, lastax.colNum + 1
        if shape[0] == nrow and shape[1] == ncol:
            axes = numpy.array(fig.axes).reshape(shape)
        else:
            fig.clear()
            # fig.set_size_inches(figsize)
            axes = fig.subplots(nrow, ncol, squeeze=False, sharex=True, sharey=True)

    for icol in range(ncol):
        hcol = h
        coltitle = None
        if col:
            vcol = col_identifiers[icol]
            hcol = h.integrate(col, vcol)
            coltitle = str(vcol)
            if isinstance(vcol, Interval) and vcol.label is None:
                coltitle = "%s ∈ %s" % (h.axis(col).label, coltitle)
        for irow in range(nrow):
            ax = axes[irow, icol]
            hplot = hcol
            rowtitle = None
            if row:
                vrow = row_identifiers[irow]
                hplot = hcol.integrate(row, vrow)
                rowtitle = str(vrow)
                if isinstance(vrow, Interval) and vrow.label is None:
                    rowtitle = "%s ∈ %s" % (h.axis(row).label, rowtitle)

            plot1d(hplot, ax=ax, overlay=overlay, **plot_opts)
            if row is not None and col is not None:
                ax.set_title("%s, %s" % (rowtitle, coltitle))
            elif row is not None:
                ax.set_title(rowtitle)
            elif col is not None:
                ax.set_title(coltitle)

    for ax in axes.flatten():
        ax.autoscale(axis="y")
        ax.set_ylim(0, None)

    return axes

def isnotebook():
    try:
        shell = get_ipython().__class__.__name__
        if shell == "ZMQInteractiveShell":
            return True  # Jupyter notebook or qtconsole
        elif shell == "TerminalInteractiveShell":
            return False  # Terminal running IPython
        else:
            return False  # Other type (?)
    except NameError:
        return False

def bokeh_plot(histo, jup_url="http://127.0.0.1:8889"):
    if not isnotebook():
        raise NotImplementedError("Only usable in jupyter notebook")
    import bokeh.plotting.figure as bk_figure
    from bokeh.io import show
    from bokeh import palettes
    from bokeh.layouts import row, column
    from bokeh.models import ColumnDataSource
    from bokeh.models.widgets import RadioButtonGroup, CheckboxButtonGroup
    from bokeh.models.widgets import RangeSlider, Div
    from bokeh.io import output_notebook  # enables plot interface in J notebook

    # init bokeh

    from bokeh.application import Application
    from bokeh.application.handlers import FunctionHandler

    from bokeh.core.validation import silence
    from bokeh.core.validation.warnings import EMPTY_LAYOUT

    silence(EMPTY_LAYOUT, True)

    output_notebook()

    # Set up widgets
    cfg_labels = ["Ghost"]
    wi_config = CheckboxButtonGroup(labels=cfg_labels, active=[0])
    wi_dense_select = RadioButtonGroup(
        labels=[ax.name for ax in histo.dense_axes()], active=0
    )
    wi_sparse_select = RadioButtonGroup(
        labels=[ax.name for ax in histo.sparse_axes()], active=0
    )

    # Dense widgets
    sliders = {}
    for ax in histo.dense_axes():
        edge_vals = (histo.axis(ax.name).edges()[0], histo.axis(ax.name).edges()[-1])
        _smallest_bin = numpy.min(numpy.diff(histo.axis(ax.name).edges()))
        sliders[ax.name] = RangeSlider(
            title=ax.name,
            value=edge_vals,
            start=edge_vals[0],
            end=edge_vals[1],
            step=_smallest_bin,
            name=ax.name,
        )

    # Cat widgets
    togglers = {}
    for ax in histo.sparse_axes():
        togglers[ax.name] = CheckboxButtonGroup(
            labels=[i.name for i in ax.identifiers()], active=[0], name=ax.name
        )

    # Toggles for all widgets
    configers = {}
    for ax in histo.sparse_axes():
        configers[ax.name] = CheckboxButtonGroup(
            labels=["Display", "Ghost"], active=[0, 1], name=ax.name
        )
    for ax in histo.dense_axes():
        configers[ax.name] = CheckboxButtonGroup(
            labels=["Display"], active=[0], name=ax.name
        )

    # Figure
    fig = bk_figure(
        title="1D Projection",
        plot_width=500,
        plot_height=500,
        min_border=20,
        toolbar_location=None,
    )
    fig.yaxis.axis_label = "N"
    fig.xaxis.axis_label = "Quantity"

    # Iterate over possible overlays
    _max_idents = 0  # Max number of simultaneou histograms
    for ax in histo.sparse_axes():
        _max_idents = max(_max_idents, len([i.name for i in ax.identifiers()]))

    # Data source list
    sources = []
    sources_ghost = []
    for i in range(_max_idents):
        sources.append(ColumnDataSource(dict(left=[], top=[], right=[], bottom=[])))
        sources_ghost.append(
            ColumnDataSource(dict(left=[], top=[], right=[], bottom=[]))
        )

    # Hist list
    hists = []
    hists_ghost = []
    for i in range(_max_idents):
        if _max_idents < 10:
            _color = palettes.Category10[min(max(3, _max_idents), 10)][i]
        else:
            _color = palettes.magma(_max_idents)[i]
        hists.append(
            fig.quad(
                left="left",
                right="right",
                top="top",
                bottom="bottom",
                source=sources[i],
                alpha=0.9,
                color=_color,
            )
        )
        hists_ghost.append(
            fig.quad(
                left="left",
                right="right",
                top="top",
                bottom="bottom",
                source=sources_ghost[i],
                alpha=0.05,
                color=_color,
            )
        )

    def update_data(attrname, old, new):
        sparse_active = wi_sparse_select.active
        sparse_name = [ax.name for ax in histo.sparse_axes()][sparse_active]
        sparse_other = [ax.name for ax in histo.sparse_axes() if ax.name != sparse_name]

        dense_active = wi_dense_select.active
        dense_name = [ax.name for ax in histo.dense_axes()][dense_active]
        dense_other = [ax.name for ax in histo.dense_axes() if ax.name != dense_name]

        # Apply cuts in projections
        _h = histo.copy()
        for proj_ax in sparse_other:
            _idents = histo.axis(proj_ax).identifiers()
            _labels = [ident.name for ident in _idents]
            if 0 in configers[proj_ax].active:
                _h = _h.integrate(
                    proj_ax, [_labels[i] for i in togglers[proj_ax].active]
                )
            else:
                _h = _h.integrate(proj_ax)

        for proj_ax in dense_other:
            _h = _h.integrate(
                proj_ax, slice(sliders[proj_ax].value[0], sliders[proj_ax].value[1])
            )

        for cat_ix in range(_max_idents):
            # Update histo for each toggled overlay
            if cat_ix in togglers[sparse_name].active:
                cat_value = histo.axis(sparse_name).identifiers()[cat_ix]
                h1d = _h.integrate(sparse_name, cat_value)

                # Get shown histogram
                values = h1d.project(dense_name).values()
                if values != {}:
                    h = values[()]
                    bins = h1d.axis(dense_name).edges()

                    # Apply cuts on shown axis
                    bin_los = bins[:-1][bins[:-1] > sliders[dense_name].value[0]]
                    bin_his = bins[1:][bins[1:] < sliders[dense_name].value[1]]
                    new_bins = numpy.intersect1d(bin_los, bin_his)
                    bin_ixs = numpy.searchsorted(bins, new_bins)[:-1]
                    h = h[bin_ixs]

                    sources[cat_ix].data = dict(
                        left=new_bins[:-1],
                        right=new_bins[1:],
                        top=h,
                        bottom=numpy.zeros_like(h),
                    )
                else:
                    sources[cat_ix].data = dict(left=[], right=[], top=[], bottom=[])

                # Add ghosts
                if 0 in wi_config.active:
                    h1d = histo.integrate(sparse_name, cat_value)
                    for proj_ax in sparse_other:
                        _idents = histo.axis(proj_ax).identifiers()
                        _labels = [ident.name for ident in _idents]
                        if 1 not in configers[proj_ax].active:
                            h1d = h1d.integrate(
                                proj_ax, [_labels[i] for i in togglers[proj_ax].active]
                            )
                        else:
                            h1d = h1d.integrate(proj_ax)
                    values = h1d.project(dense_name).values()
                    if values != {}:
                        h = h1d.project(dense_name).values()[()]
                        bins = h1d.axis(dense_name).edges()
                        sources_ghost[cat_ix].data = dict(
                            left=bins[:-1],
                            right=bins[1:],
                            top=h,
                            bottom=numpy.zeros_like(h),
                        )
                    else:
                        sources_ghost[cat_ix].data = dict(
                            left=[], right=[], top=[], bottom=[]
                        )
            else:
                sources[cat_ix].data = dict(left=[], right=[], top=[], bottom=[])
                sources_ghost[cat_ix].data = dict(left=[], right=[], top=[], bottom=[])

        # Cosmetics
        fig.xaxis.axis_label = dense_name

    for name, slider in sliders.items():
        slider.on_change("value", update_data)
    for name, toggler in togglers.items():
        toggler.on_change("active", update_data)
    for name, configer in configers.items():
        configer.on_change("active", update_data)
    # Button
    for w in [wi_dense_select, wi_sparse_select, wi_config]:
        w.on_change("active", update_data)

    from bokeh.models.widgets import Panel, Tabs

    layout = row(
        fig,
        column(
            Div(
                text="<b>Overlay Axis:</b>",
                style={"font-size": "100%", "color": "black"},
            ),
            wi_sparse_select,
            Div(
                text="<b>Plot Axis:</b>", style={"font-size": "100%", "color": "black"}
            ),
            wi_dense_select,
            Div(
                text="<b>Categorical Cuts:</b>",
                style={"font-size": "100%", "color": "black"},
            ),
            *[toggler for name, toggler in togglers.items()],
            Div(
                text="<b>Dense Cuts:</b>", style={"font-size": "100%", "color": "black"}
            ),
            *[slider for name, slider in sliders.items()]
        ),
    )

    # Config prep
    incl_lists = [[], [], []]
    for i, key in enumerate(list(configers.keys())):
        incl_lists[i // max(5, len(list(configers.keys())) / 3)].append(
            Div(
                text="<b>{}:</b>".format(key),
                style={"font-size": "70%", "color": "black"},
            )
        )
        incl_lists[i // max(5, len(list(configers.keys())) / 3)].append(configers[key])

    layout_cfgs = column(
        row(
            column(
                Div(
                    text="<b>Configs:</b>",
                    style={"font-size": "100%", "color": "black"},
                ),
                wi_config,
            )
        ),
        Div(
            text="<b>Axis togglers:</b>", style={"font-size": "100%", "color": "black"}
        ),
        row(
            column(incl_lists[0]),
            column(incl_lists[1]),
            column(incl_lists[2]),
        ),
    )

    # Update active buttons
    def update_layout(attrname, old, new):
        active_axes = [None]
        for name, wi in configers.items():
            if 0 in wi.active:
                active_axes.append(name)
        for child in layout.children[1].children:
            if child.name not in active_axes:
                child.visible = False
            else:
                child.visible = True

    for name, configer in configers.items():
        configer.on_change("active", update_layout)

    tab1 = Panel(child=layout, title="Projection")
    tab2 = Panel(child=layout_cfgs, title="Configs")
    tabs = Tabs(tabs=[tab1, tab2])

    def modify_doc(doc):
        doc.add_root(row(tabs, width=800))
        doc.title = "Sliders"

    handler = FunctionHandler(modify_doc)
    app = Application(handler)

    show(app, notebook_url=jup_url)
    update_data("", "", "")

ext: py

max_stars_repo_name: bfis/coffea

max_issues_repo_name: bfis/coffea

max_forks_repo_name: bfis/coffea

reason: import numpy,import scipy


text:
import io
import json
from pathlib import Path
from random import choice

import numpy as np
import simpleaudio as sa # sdl2_mixer non-functional for me so I resorted to this --salt-die
from PIL import Image

from kivy.app import App
from kivy.clock import Clock
from kivy.uix.widget import Widget
from kivy.graphics import Color, Rectangle

GRAVITY = .02
FRICTION = .9
DISLODGE_VELOCITY = 1e-3
MAX_VELOCITY = .2

PEBBLE_COUNT = 7e3  # per layer.
PEBBLE_IMAGE_SCALE = .75

CHISEL_RADIUS = 6e-4
MIN_POWER = 1e-5
CHISEL_POWER = 100

BACKGROUND = str(Path('assets', 'img', 'background.png'))
SOUND = tuple(str(Path('assets', 'sounds', f'00{i}.wav')) for i in range(1, 5))

def get_image_and_aspect(file):
    """
    Returns image and the correct ratio of pebbles per row and column from PEBBLE_COUNT and
    image height and width.
    """
    with Image.open(file) as image:
        w, h = image.size
        image = np.frombuffer(image.tobytes(), dtype=np.uint8)
    image = image.reshape((h, w, 4))

    pebbles_per_row = (PEBBLE_COUNT * w / h)**.5
    pebbles_per_column = pebbles_per_row * h / w

    return image, int(pebbles_per_row), int(pebbles_per_column)

PEBBLE_IMAGES = (Path("assets", "img", "boulder", f"{i}.png") for i in range(4))
PEBBLE_IMAGES = tuple(get_image_and_aspect(image) for image in PEBBLE_IMAGES)
CURRENT_IMAGE = list(choice(PEBBLE_IMAGES))

def pebble_setup():
    """
    Determines initial pebble color and placement from an image's non-transparent pixels.
    """
    image, pebbles_per_row, pebbles_per_column = CURRENT_IMAGE
    x_scale, y_scale = 1 / pebbles_per_row, 1 / pebbles_per_column
    x_offset, y_offset = (1 - PEBBLE_IMAGE_SCALE) / 2, .1  # Lower-left corner offset of image.
    h, w, _ = image.shape

    for x in range(pebbles_per_row):
        x = x_scale * x
        for y in range(pebbles_per_column):
            y = y_scale * y
            sample_loc = int(y * h), int(x * w)
            r, g, b, a = image[sample_loc]
            if not a:
                continue
            pebble_x = x * PEBBLE_IMAGE_SCALE + x_offset
            pebble_y = (1 - y) * PEBBLE_IMAGE_SCALE + y_offset
            normalized_color = r / 255, g / 255, b / 255, a / 255
            yield pebble_x, pebble_y, normalized_color

def is_dislodged(velocity):
    """
    Return False if velocity isn't enough to dislodge a pebble, else return the clipped
    velocity vector.
    """
    x, y = velocity
    magnitude = (x**2 + y**2)**.5
    if magnitude < DISLODGE_VELOCITY:
        return False
    if magnitude > MAX_VELOCITY:
        x *= MAX_VELOCITY / magnitude
        y *= MAX_VELOCITY / magnitude
    return x, y

class Pebble:
    """
    This handles physics for dislodged pebbles. Deletes itself after pebbles reach the floor.
    """
    def __init__(self, index, stone, x, y, z, velocity):
        self.index = index
        self.stone = stone
        self.pos = stone.positions[index]
        self.pixel = stone.pixels[index]
        self.x, self.y, self.z = x, y, z
        self.velocity = velocity
        self.update = Clock.schedule_interval(self.step, 0)

    def step(self, dt):
        """Gravity Physics"""
        vx, vy = self.velocity
        vx *= FRICTION
        vy *= FRICTION
        vy -= GRAVITY

        # Bounce off walls
        x, y = self.x, self.y
        if not 0 < x < 1:
            vx *= -1
        if y > 1:
            vy *= -1

        self.velocity = vx, vy
        self.x, self.y = x + vx, max(0, y + vy)

        self.update_canvas()

    def update_canvas(self):
        x, y = self.pos[:-1] = self.x, self.y
        stone = self.stone
        scaled_x, scaled_y = x * stone.width, y * stone.height
        self.pixel.pos = scaled_x, scaled_y

        if not self.y:
            self.update.cancel()
            del stone.pebbles[self.index]  # Remove reference // kill this object

class Chisel(Widget):
    """
    Handles collision detection between pebbles and the hammer.  Creates Pebbles on collision.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.sounds = tuple(sa.WaveObject.from_wave_file(sound) for sound in SOUND)
        self.setup_canvas()
        self.resize_event = Clock.schedule_once(lambda dt: None, 0)
        self.bind(size=self._delayed_resize, pos=self._delayed_resize)

    def get_pebble_size(self):
        """Calculate the correct pebble size so we have no gaps in our stone."""
        scaled_w = PEBBLE_IMAGE_SCALE * self.width
        scaled_h = PEBBLE_IMAGE_SCALE * self.height
        _, pebbles_per_row, pebbles_per_column = CURRENT_IMAGE
        return scaled_w / pebbles_per_row, scaled_h / pebbles_per_column

    def setup_canvas(self):
        self.pebbles = {}
        self.positions = []
        self.rgba = []
        self.colors = []
        self.pixels = []
        self.pebble_size = size = self.get_pebble_size()

        with self.canvas:
            self.background_color = Color(1, 1, 1, 1)
            self.background = Rectangle(pos=self.pos, size=self.size, source=BACKGROUND)

            for depth, color_scale in enumerate((.4, .6, 1)):
                for x, y, (r, g, b, a) in pebble_setup():
                    scaled_x = x * self.width
                    scaled_y = y * self.height
                    self.positions.append([x, y, depth])

                    color = color_scale * r, color_scale * g, color_scale * b, a
                    self.rgba.append(color)
                    self.colors.append(Color(*color))
                    self.pixels.append(Rectangle(pos=(scaled_x, scaled_y), size=size))

        self.background.texture.mag_filter = 'nearest'

    def _delayed_resize(self, *args):
        self.resize_event.cancel()
        self.resize_event = Clock.schedule_once(lambda dt: self.resize(*args), .3)

    def resize(self, *args):
        self.background.pos = self.pos
        self.background.size = self.size
        self.pebble_size = size = self.get_pebble_size()
        for pixel, (x, y, _) in zip(self.pixels, self.positions):
            scaled_x = x * self.width
            scaled_y = y * self.height
            pixel.pos = scaled_x, scaled_y
            pixel.size = size

    def poke_power(self, tx, ty, touch_velocity, pebble_x, pebble_y):
        """
        Returns the force vector of a poke.
        """
        dx, dy = pebble_x - tx, pebble_y - ty
        distance = dx**2 + dy**2

        if distance > CHISEL_RADIUS:
            return 0, 0
        if not distance:
            distance = 1e-4

        power = max(CHISEL_POWER * touch_velocity, MIN_POWER) / distance
        return power * dx, power * dy

    def poke(self, touch):
        """
        Apply a poke to each pebble ignoring pebbles that are below other pebbles.
        """
        tx, ty = touch.spos
        tdx, tdy = touch.dsx, touch.dsy
        touch_velocity = tdx**2 + tdy**2
        dislodged = {}

        for i, (x, y, z) in enumerate(self.positions):
            velocity = is_dislodged(self.poke_power(tx, ty, touch_velocity, x, y))
            if velocity and ((x, y) not in dislodged or dislodged[x, y][0] < z):
                dislodged[x, y] = (z, i, velocity)

        for (x, y), (z, i, velocity) in dislodged.items():
            self.pebbles[i] = Pebble(i, self, x, y, z, velocity)

    def on_touch_down(self, touch):
        self.poke(touch)
        choice(self.sounds).play()
        return True

    def on_touch_move(self, touch):
        self.poke(touch)
        return True

    def reset(self):
        CURRENT_IMAGE[:] = list(choice(PEBBLE_IMAGES))
        self.canvas.clear()
        self.setup_canvas()

    def save(self, path_to_file):
        _, pebbles_per_row, pebbles_per_column = CURRENT_IMAGE
        positions = []
        colors = []
        for (x, y, z), color in zip(self.positions, self.rgba):
            if y:
                positions.append([x, y, z])
                colors.append(color)

        pebble_dict = {'positions': positions,
                       'colors': colors,
                       'aspect': (pebbles_per_row, pebbles_per_column)}
        with open(path_to_file, 'w') as file:
            json.dump(pebble_dict, file)

    def load(self, path_to_file):
        with open(path_to_file, 'r') as file:
            pebble_dict = json.load(file)

        CURRENT_IMAGE[1:] = pebble_dict['aspect']

        self.pebbles = {}
        self.positions = pebble_dict['positions']
        self.rgba = pebble_dict['colors']
        self.colors = []
        self.pixels = []
        self.pebble_size = size = self.get_pebble_size()

        self.canvas.clear()
        with self.canvas:
            self.background_color = Color(1, 1, 1, 1)
            self.background = Rectangle(pos=self.pos, size=self.size, source=BACKGROUND)
            for (x, y, _), color in zip(self.positions, self.rgba):
                self.colors.append(Color(*color))
                scaled_x = x * self.width
                scaled_y = y * self.height
                self.pixels.append(Rectangle(pos=(scaled_x, scaled_y), size=size))
        self.background.texture.mag_filter = 'nearest'

    def export_png(self, path_to_file, transparent=False):
        colors = []  # We won't save pebbles on the floor.
        for color, pixel in zip(self.colors, self.pixels):
            _, y = pixel.pos
            if not y:
                colors.append((color, color.a))
                color.a = 0
        if transparent:
            self.background_color.a = 0

        buffer = io.BytesIO()
        self.export_as_image().save(buffer, fmt="png")

        with open(path_to_file, "wb") as file:
            file.write(buffer.getvalue())

        for color, alpha in colors:
            color.a = alpha
        self.background_color.a = 1

if __name__ == '__main__':
    class ChiselApp(App):
        def build(self):
            return Chisel()
    ChiselApp().run()

reason: import numpy


text:
import numpy as np

def phi_point_sigma(collocation,sigma,point):

    PI = 3.1415926536

    r = np.sqrt(((point[0] - collocation[0]))**2+((point[1] - collocation[1]))**2
                  +((point[2] - collocation[2]))**2)

    phi = sigma/(2*PI)*np.log(r)

    return phi

def velocity_point_sigma(collocation,sigma,point):

    PI = 3.1415926536

    r = np.sqrt(((point[0] - collocation[0]))**2+((point[1] - collocation[1]))**2
               +((point[2] - collocation[2]))**2)

    d = 0.001

    velocity_u = sigma/(2*PI*r)*(point[0] - collocation[0])/r*(1-np.exp(-1.0*(r/d)**3))
    velocity_v = sigma/(2*PI*r)*(point[1] - collocation[1])/r*(1-np.exp(-1.0*(r/d)**3))
    velocity_w = sigma/(2*PI*r)*(point[2] - collocation[2])/r*(1-np.exp(-1.0*(r/d)**3))

    velocity = np.array([velocity_u, velocity_v, velocity_w])

    return velocity


text:
# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the BSD+Patents license found in the
# LICENSE file in the root directory of this source tree.

#! /usr/bin/env python2

import numpy as np

import faiss
import unittest

class TestClustering(unittest.TestCase):

    def test_clustering(self):
        d = 64
        n = 1000
        rs = np.random.RandomState(123)
        x = rs.uniform(size=(n, d)).astype('float32')

        x *= 10

        km = faiss.Kmeans(d, 32, niter=10)
        err32 = km.train(x)

        # check that objective is decreasing
        prev = 1e50
        for o in km.obj:
            self.assertGreater(prev, o)
            prev = o

        km = faiss.Kmeans(d, 64, niter=10)
        err64 = km.train(x)

        # check that 64 centroids give a lower quantization error than 32
        self.assertGreater(err32, err64)

        km = faiss.Kmeans(d, 32, niter=10, int_centroids=True)
        err_int = km.train(x)

        # check that integer centoids are not as good as float ones
        self.assertGreater(err_int, err32)
        self.assertTrue(np.all(km.centroids == np.floor(km.centroids)))

    def test_nasty_clustering(self):
        d = 2
        rs = np.random.RandomState(123)
        x = np.zeros((100, d), dtype='float32')
        for i in range(5):
            x[i * 20:i * 20 + 20] = rs.uniform(size=d)

        # we have 5 distinct points but ask for 10 centroids...
        km = faiss.Kmeans(d, 10, niter=10, verbose=True)
        km.train(x)

    def test_redo(self):
        d = 64
        n = 1000

        rs = np.random.RandomState(123)
        x = rs.uniform(size=(n, d)).astype('float32')

        clus = faiss.Clustering(d, 20)
        clus.nredo = 1
        clus.train(x, faiss.IndexFlatL2(d))
        obj1 = faiss.vector_to_array(clus.obj)

        clus = faiss.Clustering(d, 20)
        clus.nredo = 10
        clus.train(x, faiss.IndexFlatL2(d))
        obj10 = faiss.vector_to_array(clus.obj)

        self.assertGreater(obj1[-1], obj10[-1])

class TestPCA(unittest.TestCase):

    def test_pca(self):
        d = 64
        n = 1000
        np.random.seed(123)
        x = np.random.random(size=(n, d)).astype('float32')

        pca = faiss.PCAMatrix(d, 10)
        pca.train(x)
        y = pca.apply_py(x)

        # check that energy per component is decreasing
        column_norm2 = (y**2).sum(0)

        prev = 1e50
        for o in column_norm2:
            self.assertGreater(prev, o)
            prev = o

class TestProductQuantizer(unittest.TestCase):

    def test_pq(self):
        d = 64
        n = 2000
        cs = 4
        np.random.seed(123)
        x = np.random.random(size=(n, d)).astype('float32')
        pq = faiss.ProductQuantizer(d, cs, 8)
        pq.train(x)
        codes = pq.compute_codes(x)
        x2 = pq.decode(codes)
        diff = ((x - x2)**2).sum()

        # print("diff=", diff)
        # diff= 4418.0562
        self.assertGreater(5000, diff)

        pq10 = faiss.ProductQuantizer(d, cs, 10)
        assert pq10.code_size == cs * 2
        pq10.verbose = True
        pq10.cp.verbose = True
        pq10.train(x)
        codes = pq10.compute_codes(x)

        x10 = pq10.decode(codes)
        diff10 = ((x - x10)**2).sum()
        self.assertGreater(diff, diff10)

    def do_test_codec(self, nbit):
        pq = faiss.ProductQuantizer(16, 2, nbit)

        # simulate training
        rs = np.random.RandomState(123)
        centroids = rs.rand(2, 1 << nbit, 8).astype('float32')
        faiss.copy_array_to_vector(centroids.ravel(), pq.centroids)

        idx = rs.randint(1 << nbit, size=(100, 2))
        # can be encoded exactly
        x = np.hstack((
            centroids[0, idx[:, 0]],
            centroids[1, idx[:, 1]]
        ))

        # encode / decode
        codes = pq.compute_codes(x)
        xr = pq.decode(codes)
        assert np.all(xr == x)

        # encode w/ external index
        assign_index = faiss.IndexFlatL2(8)
        pq.assign_index = assign_index
        codes2 = np.empty((100, pq.code_size), dtype='uint8')
        pq.compute_codes_with_assign_index(
            faiss.swig_ptr(x), faiss.swig_ptr(codes2), 100)

        assert np.all(codes == codes2)

    def test_codec(self):
        self.do_test_codec(6)
        self.do_test_codec(10)
        # self.do_test_codec(16)

class TestRevSwigPtr(unittest.TestCase):

    def test_rev_swig_ptr(self):

        index = faiss.IndexFlatL2(4)
        xb0 = np.vstack([
            i * 10 + np.array([1, 2, 3, 4], dtype='float32')
            for i in range(5)])
        index.add(xb0)
        xb = faiss.rev_swig_ptr(index.xb.data(), 4 * 5).reshape(5, 4)
        self.assertEqual(np.abs(xb0 - xb).sum(), 0)

class TestException(unittest.TestCase):

    def test_exception(self):

        index = faiss.IndexFlatL2(10)

        a = np.zeros((5, 10), dtype='float32')
        b = np.zeros(5, dtype='int64')

        try:
            # an unsupported operation for IndexFlat
            index.add_with_ids(a, b)
        except RuntimeError as e:
            assert 'add_with_ids not implemented' in str(e)
        else:
            assert False, 'exception did not fire???'

    def test_exception_2(self):

        try:
            faiss.index_factory(12, 'IVF256,Flat,PQ8')
        except RuntimeError as e:
            assert 'could not parse' in str(e)
        else:
            assert False, 'exception did not fire???'

class TestMapLong2Long(unittest.TestCase):

    def test_maplong2long(self):
        keys = np.array([13, 45, 67])
        vals = np.array([3, 8, 2])

        m = faiss.MapLong2Long()
        m.add(keys, vals)

        assert np.all(m.search_multiple(keys) == vals)

        assert m.search(12343) == -1

class TestOrthognalReconstruct(unittest.TestCase):

    def test_recons_orthonormal(self):
        lt = faiss.LinearTransform(20, 10, True)
        rs = np.random.RandomState(10)
        A, _ = np.linalg.qr(rs.randn(20, 20))
        A = A[:10].astype('float32')
        faiss.copy_array_to_vector(A.ravel(), lt.A)
        faiss.copy_array_to_vector(rs.randn(10).astype('float32'), lt.b)

        lt.set_is_orthonormal()
        lt.is_trained = True
        assert lt.is_orthonormal

        x = rs.rand(30, 20).astype('float32')
        xt = lt.apply_py(x)
        xtt = lt.reverse_transform(xt)
        xttt = lt.apply_py(xtt)

        err = ((xt - xttt)**2).sum()

        self.assertGreater(1e-5, err)

    def test_recons_orthogona_impossible(self):
        lt = faiss.LinearTransform(20, 10, True)
        rs = np.random.RandomState(10)
        A = rs.randn(10 * 20).astype('float32')
        faiss.copy_array_to_vector(A.ravel(), lt.A)
        faiss.copy_array_to_vector(rs.randn(10).astype('float32'), lt.b)
        lt.is_trained = True

        lt.set_is_orthonormal()
        assert not lt.is_orthonormal

        x = rs.rand(30, 20).astype('float32')
        xt = lt.apply_py(x)
        try:
            lt.reverse_transform(xt)
        except Exception:
            pass
        else:
            self.assertFalse('should do an exception')

class TestMAdd(unittest.TestCase):

    def test_1(self):
        # try with dimensions that are multiples of 16 or not
        rs = np.random.RandomState(123)
        swig_ptr = faiss.swig_ptr
        for dim in 16, 32, 20, 25:
            for repeat in 1, 2, 3, 4, 5:
                a = rs.rand(dim).astype('float32')
                b = rs.rand(dim).astype('float32')
                c = np.zeros(dim, dtype='float32')
                bf = rs.uniform(5.0) - 2.5
                idx = faiss.fvec_madd_and_argmin(
                    dim, swig_ptr(a), bf, swig_ptr(b),
                    swig_ptr(c))
                ref_c = a + b * bf
                assert np.abs(c - ref_c).max() < 1e-5
                assert idx == ref_c.argmin()

class TestNyFuncs(unittest.TestCase):

    def test_l2(self):
        rs = np.random.RandomState(123)
        swig_ptr = faiss.swig_ptr
        for d in 1, 2, 4, 8, 12, 16:
            x = rs.rand(d).astype('float32')
            for ny in 128, 129, 130:
                print("d=%d ny=%d" % (d, ny))
                y = rs.rand(ny, d).astype('float32')
                ref = ((x - y) ** 2).sum(1)
                new = np.zeros(ny, dtype='float32')
                faiss.fvec_L2sqr_ny(swig_ptr(new), swig_ptr(x),
                                    swig_ptr(y), d, ny)
                assert np.abs(ref - new).max() < 1e-4

    def test_IP(self):
        # this one is not optimized with SIMD but just in case
        rs = np.random.RandomState(123)
        swig_ptr = faiss.swig_ptr
        for d in 1, 2, 4, 8, 12, 16:
            x = rs.rand(d).astype('float32')
            for ny in 128, 129, 130:
                print("d=%d ny=%d" % (d, ny))
                y = rs.rand(ny, d).astype('float32')
                ref = (x * y).sum(1)
                new = np.zeros(ny, dtype='float32')
                faiss.fvec_inner_products_ny(
                    swig_ptr(new), swig_ptr(x), swig_ptr(y), d, ny)
                assert np.abs(ref - new).max() < 1e-4

class TestMatrixStats(unittest.TestCase):

    def test_0s(self):
        rs = np.random.RandomState(123)
        m = rs.rand(40, 20).astype('float32')
        m[5:10] = 0
        comments = faiss.MatrixStats(m).comments
        print(comments)
        assert 'has 5 copies' in comments
        assert '5 null vectors' in comments

    def test_copies(self):
        rs = np.random.RandomState(123)
        m = rs.rand(40, 20).astype('float32')
        m[::2] = m[1::2]
        comments = faiss.MatrixStats(m).comments
        print(comments)
        assert '20 vectors are distinct' in comments

    def test_dead_dims(self):
        rs = np.random.RandomState(123)
        m = rs.rand(40, 20).astype('float32')
        m[:, 5:10] = 0
        comments = faiss.MatrixStats(m).comments
        print(comments)
        assert '5 dimensions are constant' in comments

    def test_rogue_means(self):
        rs = np.random.RandomState(123)
        m = rs.rand(40, 20).astype('float32')
        m[:, 5:10] += 12345
        comments = faiss.MatrixStats(m).comments
        print(comments)
        assert '5 dimensions are too large wrt. their variance' in comments

    def test_normalized(self):
        rs = np.random.RandomState(123)
        m = rs.rand(40, 20).astype('float32')
        faiss.normalize_L2(m)
        comments = faiss.MatrixStats(m).comments
        print(comments)
        assert 'vectors are normalized' in comments

class TestScalarQuantizer(unittest.TestCase):

    def test_8bit_equiv(self):
        rs = np.random.RandomState(123)
        for it in range(20):
            for d in 13, 16, 24:
                x = np.floor(rs.rand(5, d) * 256).astype('float32')
                x[0] = 0
                x[1] = 255

                # make sure to test extreme cases
                x[2, 0] = 0
                x[3, 0] = 255
                x[2, 1] = 255
                x[3, 1] = 0

                ref_index = faiss.IndexScalarQuantizer(d, faiss.ScalarQuantizer.QT_8bit)
                ref_index.train(x[:2])
                ref_index.add(x[2:3])

                index = faiss.IndexScalarQuantizer(
                    d, faiss.ScalarQuantizer.QT_8bit_direct)
                assert index.is_trained
                index.add(x[2:3])

                assert np.all(
                    faiss.vector_to_array(ref_index.codes) ==
                    faiss.vector_to_array(index.codes))

                # Note that distances are not the same because ref_index
                # reconstructs x as x + 0.5
                D, I = index.search(x[3:], 1)

                # assert D[0, 0] == Dref[0, 0]
                print(D[0, 0], ((x[3] - x[2]) ** 2).sum())
                assert D[0, 0] == ((x[3] - x[2]) ** 2).sum()

if __name__ == '__main__':
    unittest.main()


text:
import soundata
import numpy as np
import matplotlib.pyplot as plt
from micarraylib.arraycoords.array_shapes_utils import _polar2cart

class Aggregate:
    """
    The generic dataset aggregator class

    Args:
        micarray_dataset_list (list): a list with the
            micarray.Dataset objects to aggregate
        fs (float): the sampling rate to use across
            datasets in the aggregate
    """

    def __init__(self, micarray_dataset_list, fs):
        self.fs = fs
        self.datasets = {dataset.name: dataset for dataset in micarray_dataset_list}
        for dataset in self.datasets.values():
            dataset.fs = fs

class Dataset:
    """
    The generic Dataset class

    Args:
        name (str): the dataset name (must be one of
            soundata dataset names)
        fs (int): the dataset's global sampling rate
            (indicated in the soundata loader)
        array_format (dict): a dictionary specifying
            the format (A or B) of the original audio files
            for each microphone array in the dataset
        capsule_coords (dict): a dictionary specifying
            each of the array's capsule polar coordinates,
            with the array's physical center at the
            origin of the 3D space.
        download (bool): whether the data needs to
            be downloaded by soundata (True by default)
        data_home (str): path to directory where the
            data is (directories must follow
            the soundata index struct)

    Attributes:
        name (str): the dataset name (must be one of
            soundata dataset names)
        dataset (soundata.Dataset)
        fs (int): the dataset's global sampling rate
            (indicated in the soundata loader)
        array_format (dict): a dictionary specifying
            the format (A or B) of the original audio files
            for each microphone array in the dataset
        array_names (list): list with all the arrays
            in the dataset
        capsule_coords (dict): a dictionary specifying
            each of the array's capsule polar coordinates,
            with the array's physical center at the
            origin of the 3D space.
        array_capsules (dict):
        data_home (str): path to directory where the
            data is (directories inside data_home must
            follow the soundata index struct)
    """

    def __init__(
        self,
        name,
        fs,
        array_format,
        capsule_coords,
        download=True,
        data_home=None,
        partial_download=None,
        force_overwrite=False,
        cleanup=False,
    ):

        if download == False and data_home == None:
            raise ValueError(
                "You must specify the directory with the data in data_home if you do not want to download it."
            )
        self.name = name
        self.dataset = _initialize(
            name, data_home, download, partial_download, force_overwrite, cleanup
        )
        self.fs = fs
        self.array_names = list(capsule_coords.keys())
        self.array_format = array_format
        self.array_capsules = {
            a: list(capsule_coords[a].keys()) for a in capsule_coords.keys()
        }
        self.capsule_coords = capsule_coords
        self.data_home = data_home

    def get_capsule_coords_numpy(self, micarray):
        """
        returns the capsule coordinates in
        polar form

        Args:
            micarray (str): the name of the
                microphone array that we are
                getting coordinates for

        Returns:
            tuple with
                1) numpy array of shape (n,3)
                2) capsule coordinate names (list)
        """

        if micarray not in self.array_names:
            raise ValueError(
                "micarray is {}, but it should be one of {}".format(
                    micarray, ", ".join(self.array_names)
                )
            )
        capsule_coords = np.array([c for c in self.capsule_coords[micarray].values()])
        capsule_names = [c for c in self.capsule_coords[micarray].keys()]
        return capsule_coords, capsule_names

    def plot_micarray(self, micarray, show=True):
        """
        returns the capsule coordinates in
        polar form

        Args:
            micarray (str): the name of the
                microphone array that we are
                getting coordinates for

        """

        if micarray not in self.array_names:
            raise ValueError(
                "micarray is {}, but it should be one of {}".format(
                    micarray, ", ".join(self.array_names)
                )
            )
        capsule_coords, capsule_names = self.get_capsule_coords_numpy(micarray)
        capsule_coords_dict = _polar2cart(
            {c: capsule_coords[i] for i, c in enumerate(capsule_names)}, "radians"
        )
        fig = plt.figure()
        ax = plt.axes(projection="3d")
        for n, m in capsule_coords_dict.items():
            ax.scatter(m[0], m[1], m[2], color="b", label=micarray)
            ax.text(m[0], m[1], m[2], "%s" % (str(n)), size=20, zorder=1, color="k")

        ax.set_xlabel("x (meters)")
        ax.set_ylabel("y (meters)")
        ax.set_zlabel("z (meters)")
        if show:
            plt.show()

def _initialize(
    name,
    data_home=None,
    download=True,
    partial_download=None,
    force_overwrite=False,
    cleanup=False,
):
    """
    Initializes and downloads the dataset
    using the soundata API

    Args:
        data_home (str): points to the directory with
            the data
        download (bool): True by default. soundata
            will skip this step if the files already exist
    """
    dataset = soundata.initialize(name, data_home)
    if download:
        dataset.download(partial_download, force_overwrite, cleanup)
    return dataset


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Oct 22 16:35:33 2019

@author: nico
"""
import sys
sys.path.append('/home/nico/Documentos/facultad/6to_nivel/pds/git/pdstestbench')
import os
import matplotlib.pyplot as plt
import numpy as np
from scipy.fftpack import fft
from pdsmodulos.signals import FFT
from pdsmodulos.signals import windows as wds
from pdsmodulos.signals import signals as sig

#import seaborn as sns

os.system ("clear") # limpia la terminal de python
plt.close("all")    #cierra todos los graficos

N  = 1000 # muestras
fs = 2*np.pi # Hz
df = fs / N
a0 = 1 # Volts
p0 = 0 # radianes
f0 = np.pi / 2
dif= -40  #Diferencia en dB de Amplitude
a1 = a0 * 10**(dif/20)
f1 = f0 + 10 * df

#%% Señales ej 2a
tt = np.linspace(0, (N-1)/fs, N)
freq = np.linspace(0, (N-1)*df, N) / fs

signal = a0 * np.sin(2*np.pi*f0*tt) + a1 * np.sin(2*np.pi*f1*tt)
mod_signal=np.abs(np.fft.fft(signal))*2/N
mod_signal = 20 *np.log10(mod_signal)

# grafico
fig = plt.figure("Señal bitonal " , constrained_layout=True)
plt.title("Señal bitonal separada 10df")
plt.plot(freq[0:int(N/2)],mod_signal[0:int(N/2)])
plt.xlabel("Frecuencia normalizada [f/fs]")
plt.ylabel("Magnitud [dB]")
plt.axhline(0, color="black")
plt.axvline(0, color="black")
plt.grid()

#%% Señales ej 2b
a1 = a0 * 10**(-265/20)
signal2 = a0 * np.sin(2*np.pi*f0*tt) + a1 * np.sin(2*np.pi*f1*tt)
mod_signal2 = np.abs(np.fft.fft(signal2))*2/N
mod_signal2 = 20 *np.log10(mod_signal2)

# grafico
fig = plt.figure("Señal bitonal con a1 minimo " , constrained_layout=True)
plt.title("Señal bitonal separada 10df")
plt.plot(freq[0:int(N/2)],mod_signal2[0:int(N/2)])
plt.xlabel("Frecuencia normalizada [f/fs]")
plt.ylabel("Magnitud [dB]")
plt.axhline(0, color="black")
plt.axvline(0, color="black")
plt.grid()

#%% Genero ruido de cuantificacion
a1 = a0 * 10**(-90/20)
signal3 = a0 * np.sin(2*np.pi*f0*tt) + a1 * np.sin(2*np.pi*f1*tt)
# para que se vea mejor el ruido ya que es muy deterministico agregarle  ruido con distribuzión
#uniforme de -/- 5 LSB si lo hago despues del cuentificador de -5 a +5
n = 16 # bits de cuatizacion
LSB = a0/2**n
ruido = np.random.uniform(-5*LSB, 5*LSB, N)
signal3 = signal3 + ruido

qsignal3 = np.round(signal3*(2**n)/2) / ((2**n)/2)

e16 = signal3 - qsignal3

mod_signal3 = np.abs(fft(qsignal3))*2/N # FFT de la señal cuantificada
mod_signal3 = 20 *np.log10(mod_signal3 + np.finfo(float).eps)

#tt = np.linspace(0,(N-1)/fs, N)
#fig = plt.figure("Señal bitonal cuantificada en tiempo ")
#plt.title("Señal bitonal separada 10df")
#plt.plot(tt[:20], qsignal3[:20])
#plt.xlabel("Tiempo")
#plt.ylabel("Amplitud")
#plt.axhline(0, color="black")
#plt.axvline(0, color="black")
#plt.grid()

fig = plt.figure("Señal bitonal cuantificada ")
plt.title("Señal bitonal separada 10df")
plt.plot(freq[0:int(N/2)],mod_signal3[0:int(N/2)])
plt.xlabel("Frecuencia normalizada [f/fs]")
plt.ylabel("Magnitud [dB]")
plt.axhline(0, color="black")
plt.axvline(0, color="black")
plt.grid()

max_stars_repo_name: NicolasFerragamo/pdstestbench

max_issues_repo_name: NicolasFerragamo/pdstestbench

max_forks_repo_name: NicolasFerragamo/pdstestbench

reason: import numpy,from scipy