text:
# This code is part of Qiskit.
#
# (C) Copyright IBM 2021.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""Different data analysis steps."""

from abc import abstractmethod
from abc import ABC
from enum import Enum
from numbers import Number
from typing import Union, Sequence, Set
from collections import defaultdict

import numpy as np
from uncertainties import unumpy as unp, ufloat

from qiskit.result.postprocess import format_counts_memory
from qiskit_experiments.data_processing.data_action import DataAction, TrainableDataAction
from qiskit_experiments.data_processing.exceptions import DataProcessorError
from qiskit_experiments.framework import Options

class AverageData(DataAction):
    """A node to average data representable as numpy arrays."""

    def __init__(self, axis: int, validate: bool = True):
        r"""Initialize a data averaging node.

        Args:
            axis: The axis along which to average.
            validate: If set to False the DataAction will not validate its input.

        Notes:
            If the input array has no standard error, then this node will compute the
            standard error of the mean, i.e. the standard deviation of the datum divided by
            :math:`\sqrt{N}` where :math:`N` is the number of data points.
            Otherwise the standard error is given by the square root of :math:`N^{-1}` times
            the sum of the squared errors.
        """
        super().__init__(validate)
        self._axis = axis

    def _format_data(self, data: np.ndarray) -> np.ndarray:
        """Format the data into numpy arrays.

        Args:
            data: A data array to format. This is a single numpy array containing
                all circuit results input to the data processor.

        Returns:
            The data that has been validated and formatted.

        Raises:
            DataProcessorError: When the specified axis does not exist in given array.
        """
        if self._validate:
            if len(data.shape) <= self._axis:
                raise DataProcessorError(
                    f"Cannot average the {len(data.shape)} dimensional "
                    f"array along axis {self._axis}."
                )

        return data

    def _process(self, data: np.ndarray) -> np.ndarray:
        """Average the data.

        Args:
            data: A data array to process. This is a single numpy array containing
                all circuit results input to the data processor.

        Returns:
             Arrays with one less dimension than the given data.
        """
        ax = self._axis

        reduced_array = np.mean(data, axis=ax)
        nominals = unp.nominal_values(reduced_array)
        errors = unp.std_devs(reduced_array)

        if np.any(np.isnan(errors)):
            # replace empty elements with SEM
            sem = np.std(unp.nominal_values(data), axis=ax) / np.sqrt(data.shape[ax])
            errors = np.where(np.isnan(errors), sem, errors)

        return unp.uarray(nominals, errors)

    def __repr__(self):
        """String representation of the node."""
        return f"{self.__class__.__name__}(validate={self._validate}, axis={self._axis})"

class MinMaxNormalize(DataAction):
    """Normalizes the data."""

    def _process(self, data: np.ndarray) -> np.ndarray:
        """Normalize the data to the interval [0, 1].

        Args:
            data: A data array to process. This is a single numpy array containing
                all circuit results input to the data processor.

        Returns:
            The normalized data.

        Notes:
            This doesn't consider the uncertainties of the minimum or maximum values.
            Input data array is just scaled by the data range.
        """
        nominals = unp.nominal_values(data)
        min_y, max_y = np.min(nominals), np.max(nominals)

        return (data - min_y) / (max_y - min_y)

class SVD(TrainableDataAction):
    """Singular Value Decomposition of averaged IQ data."""

    def __init__(self, validate: bool = True):
        """Create new action.

        Args:
            validate: If set to False the DataAction will not validate its input.
        """
        super().__init__(validate=validate)
        self._n_circs = 0
        self._n_shots = 0
        self._n_slots = 0
        self._n_iq = 0

    @classmethod
    def _default_parameters(cls) -> Options:
        """Default parameters.

        Parameters are defined for each qubit in the data and thus
        represented as an array-like.

        Trainable parameters:
            main_axes: A unit vector representing the main axis in IQ plane.
            i_means: Mean value of training data along I quadrature.
            q_means: Mean value of training data along Q quadrature.
            scales: Scaling of IQ signal.
        """
        params = super()._default_parameters()
        params.main_axes = None
        params.i_means = None
        params.q_means = None
        params.scales = None

        return params

    def _format_data(self, data: np.ndarray) -> np.ndarray:
        """Check that the IQ data is 2D and convert it to a numpy array.

        Args:
            data: A data array to format. This is a single numpy array containing
                all circuit results input to the data processor.
                This data has different dimensions depending on whether
                single-shot or averaged data is being processed.
                Single-shot data is four dimensional, i.e., ``[#circuits, #shots, #slots, 2]``,
                while averaged IQ data is three dimensional, i.e., ``[#circuits, #slots, 2]``.
                Here, ``#slots`` is the number of classical registers used in the circuit.

        Returns:
            data and any error estimate as a numpy array.

        Raises:
            DataProcessorError: If the datum does not have the correct format.
        """
        self._n_circs = 0
        self._n_shots = 0
        self._n_slots = 0
        self._n_iq = 0

        # identify shape
        try:
            # level1 single-shot data
            self._n_circs, self._n_shots, self._n_slots, self._n_iq = data.shape
        except ValueError:
            try:
                # level1 data averaged over shots
                self._n_circs, self._n_slots, self._n_iq = data.shape
            except ValueError as ex:
                raise DataProcessorError(
                    f"Data given to {self.__class__.__name__} is not likely level1 data."
                ) from ex

        if self._validate:
            if self._n_iq != 2:
                raise DataProcessorError(
                    f"IQ data given to {self.__class__.__name__} does not have two-dimensions "
                    f"(I and Q). Instead, {self._n_iq} dimensions were found."
                )

        return data

    def _process(self, data: np.ndarray) -> np.ndarray:
        """Project the IQ data onto the axis defined by an SVD and scale it.

        Args:
            data: A data array to process. This is a single numpy array containing
                all circuit results input to the data processor.

        Returns:
            A Tuple of 1D arrays of the result of the SVD and the associated error. Each entry
            is the real part of the averaged IQ data of a qubit.

        Raises:
            DataProcessorError: If the SVD has not been previously trained on data.
        """
        if not self.is_trained:
            raise DataProcessorError("SVD must be trained on data before it can be used.")

        # IQ axis is reduced by projection
        if self._n_shots == 0:
            # level1 average mode
            dims = self._n_circs, self._n_slots
        else:
            # level1 single mode
            dims = self._n_circs, self._n_shots, self._n_slots

        projected_data = np.zeros(dims, dtype=object)

        for idx in range(self._n_slots):
            scale = self.parameters.scales[idx]
            # error propagation is computed from data if any std error exists
            centered = np.array(
                [
                    data[..., idx, 0] - self.parameters.i_means[idx],
                    data[..., idx, 1] - self.parameters.q_means[idx],
                ]
            )
            projected_data[..., idx] = (self.parameters.main_axes[idx] @ centered) / scale

        return projected_data

    def train(self, data: np.ndarray):
        """Train the SVD on the given data.

        Each element of the given data will be converted to a 2D array of dimension
        n_qubits x 2. The number of qubits is inferred from the shape of the data.
        For each qubit the data is collected into an array of shape 2 x n_data_points.
        The mean of the in-phase a quadratures is subtracted before passing the data
        to numpy's svd function. The dominant axis and the scale is saved for each
        qubit so that future data points can be projected onto the axis.

        .. note::

            Before training the SVD the mean of the training data is subtracted from the
            training data to avoid large offsets in the data.
            These means can be retrieved with :attr:`.parameters.i_means` or
            :attr:`parameters.q_means` for I and Q quadrature, respectively.

        Args:
            data: A data array to be trained. This is a single numpy array containing
                all circuit results input to the data processor.
        """
        if data is None:
            return

        # TODO do not remove standard error. Currently svd is not supported.
        data = unp.nominal_values(self._format_data(data))

        main_axes = []
        scales = []
        i_means = []
        q_means = []
        for idx in range(self._n_slots):
            datums = np.vstack([datum[idx] for datum in data]).T

            # Calculate the mean of the data to recenter it in the IQ plane.
            mean_i = np.average(datums[0, :])
            mean_q = np.average(datums[1, :])
            i_means.append(mean_i)
            q_means.append(mean_q)

            datums[0, :] = datums[0, :] - mean_i
            datums[1, :] = datums[1, :] - mean_q

            mat_u, mat_s, _ = np.linalg.svd(datums)
            main_axes.append(mat_u[:, 0])
            scales.append(mat_s[0])

        self.set_parameters(
            main_axes=main_axes,
            scales=scales,
            i_means=i_means,
            q_means=q_means,
        )

class IQPart(DataAction):
    """Abstract class for IQ data post-processing."""

    def __init__(self, scale: float = 1.0, validate: bool = True):
        """
        Args:
            scale: Float with which to multiply the IQ data. Defaults to 1.0.
            validate: If set to False the DataAction will not validate its input.
        """
        self.scale = scale
        super().__init__(validate)

    @abstractmethod
    def _process(self, data: np.ndarray) -> np.ndarray:
        """Defines how the IQ point is processed.

        The last dimension of the array should correspond to [real, imaginary] part of data.

        Args:
            data: A data array to process. This is a single numpy array containing
                all circuit results input to the data processor.

        Returns:
            The data that has been processed.
        """

    def _format_data(self, data: np.ndarray) -> np.ndarray:
        """Format and validate the input.

        Args:
            data: A data array to format. This is a single numpy array containing
                all circuit results input to the data processor.

        Returns:
            The data that has been validated and formatted.

        Raises:
            DataProcessorError: When input data is not likely IQ data.
        """
        if self._validate:
            if data.shape[-1] != 2:
                raise DataProcessorError(
                    f"IQ data given to {self.__class__.__name__} must be a multi-dimensional array"
                    "of dimension [d0, d1, ..., 2] in which the last dimension "
                    "corresponds to IQ elements."
                    f"Input data contains element with length {data.shape[-1]} != 2."
                )

        return data

    def __repr__(self):
        """String representation of the node."""
        return f"{self.__class__.__name__}(validate={self._validate}, scale={self.scale})"

class ToReal(IQPart):
    """IQ data post-processing. Isolate the real part of single-shot IQ data."""

    def _process(self, data: np.ndarray) -> np.ndarray:
        """Take the real part of the IQ data.

        Args:
            data: An N-dimensional array of complex IQ point as [real, imaginary].

        Returns:
            A N-1 dimensional array, each entry is the real part of the given IQ data.
        """
        return data[..., 0] * self.scale

class ToImag(IQPart):
    """IQ data post-processing. Isolate the imaginary part of single-shot IQ data."""

    def _process(self, data: np.ndarray) -> np.ndarray:
        """Take the imaginary part of the IQ data.

        Args:
            data: An N-dimensional array of complex IQ point as [real, imaginary].

        Returns:
            A N-1 dimensional array, each entry is the imaginary part of the given IQ data.
        """
        return data[..., 1] * self.scale

class ToAbs(IQPart):
    """IQ data post-processing. Take the absolute value of the IQ point."""

    def _process(self, data: np.array) -> np.array:
        """Take the absolute value of the IQ data.

        Args:
            data: An N-dimensional array of complex IQ point as [real, imaginary].

        Returns:
            A N-1 dimensional array, each entry is the absolute value of the given IQ data.
        """
        # pylint: disable=no-member
        return unp.sqrt(data[..., 0] ** 2 + data[..., 1] ** 2) * self.scale

class CountsAction(DataAction):
    """An abstract DataAction that acts on count dictionaries."""

    @abstractmethod
    def _process(self, data: np.ndarray) -> np.ndarray:
        """Defines how the counts are processed."""

    def _format_data(self, data: np.ndarray) -> np.ndarray:
        """
        Checks that the given data has a counts format.

        Args:
            data: A data array to format. This is a single numpy array containing
                all circuit results input to the data processor.
                This is usually an object data type containing Python dictionaries of
                count data keyed on the measured bitstring.
                A count value is a discrete quantity representing the frequency of an event.
                Therefore, count values do not have an uncertainty.

        Returns:
            The ``data`` as given.

        Raises:
            DataProcessorError: If the data is not a counts dict or a list of counts dicts.
        """
        valid_count_type = int, np.integer

        if self._validate:
            for datum in data:
                if not isinstance(datum, dict):
                    raise DataProcessorError(
                        f"Data entry must be dictionary of counts, received {type(datum)}."
                    )
                for bit_str, count in datum.items():
                    if not isinstance(bit_str, str):
                        raise DataProcessorError(
                            f"Key {bit_str} is not a valid count key in {self.__class__.__name__}."
                        )
                    if not isinstance(count, valid_count_type):
                        raise DataProcessorError(
                            f"Count {bit_str} is not a valid count for {self.__class__.__name__}. "
                            "The uncertainty of probability is computed based on sampling error, "
                            "thus the count should be an error-free discrete quantity "
                            "representing the frequency of event."
                        )

        return data

class MarginalizeCounts(CountsAction):
    r"""A data action to marginalize count dictionaries.

    This data action takes a count dictionary and returns a new count dictionary that
    is marginalized over a set of specified qubits. For example, given the count
    dictionary :code:`{"010": 1, "110": 10, "100": 100}` this node will return the
    count dictionary :code:`{"10": 11, "00": 100}` when marginalized over qubit 2.

    .. note::
        This data action can be used to discard one or more qubits in the counts
        dictionary. This is, for example, useful when processing two-qubit restless
        experiments but can be used in a more general context. In composite
        experiments the counts marginalization is already done in the data container.
    """

    def __init__(self, qubits_to_keep: Set[int], validate: bool = True):
        """Initialize a counts marginalization node.

        Args:
            qubits_to_keep: A set of qubits to retain during the marginalization process.
            validate: If set to False the DataAction will not validate its input.
        """
        super().__init__(validate)
        self._qubits_to_keep = sorted(qubits_to_keep, reverse=True)

    def _process(self, data: np.ndarray) -> np.ndarray:
        """Perform counts marginalization."""
        marginalized_counts = []

        for datum in data:
            new_counts = defaultdict(int)
            for bit_str, count in datum.items():
                new_counts["".join([bit_str[::-1][idx] for idx in self._qubits_to_keep])] += count

            marginalized_counts.append(new_counts)

        return np.array(marginalized_counts)

    def __repr__(self):
        """String representation of the node."""
        options_str = ", ".join(
            [
                f"qubits_to_keep={self._qubits_to_keep}",
                f"validate={self._validate}",
            ]
        )
        return f"{self.__class__.__name__}({options_str})"

class Probability(CountsAction):
    r"""Compute the mean probability of a single measurement outcome from counts.

    This node returns the mean and standard deviation of a single measurement
    outcome probability :math:`p` estimated from the observed counts. The mean and
    variance are computed from the posterior Beta distribution
    :math:`B(\alpha_0^\prime,\alpha_1^\prime)` estimated from a Bayesian update
    of a prior Beta distribution :math:`B(\alpha_0, \alpha_1)` given the observed
    counts.

    The mean and variance of the Beta distribution :math:`B(\alpha_0, \alpha_1)` are:

    .. math::

        \text{E}[p] = \frac{\alpha_0}{\alpha_0 + \alpha_1}, \quad
        \text{Var}[p] = \frac{\text{E}[p] (1 - \text{E}[p])}{\alpha_0 + \alpha_1 + 1}

    Given a prior Beta distribution :math:`B(\alpha_0, \alpha_1)`, the posterior
    distribution for the observation of :math:`F` counts of a given
    outcome out of :math:`N` total shots is a
    :math:`B(\alpha_0^\prime,\alpha_1^\prime):math:` with

    .. math::
        \alpha_0^\prime = \alpha_0 + F, \quad
        \alpha_1^\prime = \alpha_1 + N - F.

    .. note::

        The default value for the prior distribution is *Jeffery's Prior*
        :math:`\alpha_0 = \alpha_1 = 0.5` which represents ignorance about the true
        probability value. Note that for this prior the mean probability estimate
        from a finite number of counts can never be exactly 0 or 1. The estimated
        mean and variance are given by

        .. math::

            \text{E}[p] = \frac{F + 0.5}{N + 1}, \quad
            \text{Var}[p] = \frac{\text{E}[p] (1 - \text{E}[p])}{N + 2}

        This node will deprecate standard error provided by the previous node.
    """

    def __init__(
        self,
        outcome: str,
        alpha_prior: Union[float, Sequence[float]] = 0.5,
        validate: bool = True,
    ):
        """Initialize a counts to probability data conversion.

        Args:
            outcome: The bitstring for which to return the probability and variance.
            alpha_prior: A prior Beta distribution parameter ``[`alpha0, alpha1]``.
                         If specified as float this will use the same value for
                         ``alpha0`` and``alpha1`` (Default: 0.5).
            validate: If set to False the DataAction will not validate its input.

        Raises:
            DataProcessorError: When the dimension of the prior and expected parameter vector
                do not match.
        """
        self._outcome = outcome
        if isinstance(alpha_prior, Number):
            self._alpha_prior = [alpha_prior, alpha_prior]
        else:
            if validate and len(alpha_prior) != 2:
                raise DataProcessorError(
                    "Prior for probability node must be a float or pair of floats."
                )
            self._alpha_prior = list(alpha_prior)

        super().__init__(validate)

    def _process(self, data: np.ndarray) -> np.ndarray:
        """Compute mean and standard error from the beta distribution.

        Args:
            data: A data array to process. This is a single numpy array containing
                all circuit results input to the data processor.
                This is usually an object data type containing Python dictionaries of
                count data keyed on the measured bitstring.

        Returns:
            The data that has been processed.
        """
        probabilities = np.empty(data.size, dtype=object)

        for idx, counts_dict in enumerate(data):
            shots = sum(counts_dict.values())
            freq = counts_dict.get(self._outcome, 0)
            alpha_posterior = [freq + self._alpha_prior[0], shots - freq + self._alpha_prior[1]]
            alpha_sum = sum(alpha_posterior)

            p_mean = alpha_posterior[0] / alpha_sum
            p_var = p_mean * (1 - p_mean) / (alpha_sum + 1)

            probabilities[idx] = ufloat(nominal_value=p_mean, std_dev=np.sqrt(p_var))

        return probabilities

    def __repr__(self):
        """String representation of the node."""
        options_str = ", ".join(
            [
                f"validate={self._validate}",
                f"outcome={self._outcome}",
                f"alpha_prior={self._alpha_prior}",
            ]
        )
        return f"{self.__class__.__name__}({options_str})"

class BasisExpectationValue(DataAction):
    """Compute expectation value of measured basis from probability.

    Note:
        The sign becomes P(0) -> 1, P(1) -> -1.
    """

    def _format_data(self, data: np.ndarray) -> np.ndarray:
        """Format and validate the input.

        Args:
            data: A data array to format. This is a single numpy array containing
                all circuit results input to the data processor.

        Returns:
            The data that has been validated and formatted.

        Raises:
            DataProcessorError: When input value is not in [0, 1]
        """
        if self._validate:
            if not all(0.0 <= p <= 1.0 for p in data):
                raise DataProcessorError(
                    f"Input data for node {self.__class__.__name__} is not likely probability."
                )

        return data

    def _process(self, data: np.ndarray) -> np.ndarray:
        """Compute basis eigenvalue.

        Args:
            data: A data array to process. This is a single numpy array containing
                all circuit results input to the data processor.

        Returns:
            The data that has been processed.
        """
        return 2 * (0.5 - data)

class ProjectorType(Enum):
    """Types of projectors for data dimensionality reduction."""

    SVD = SVD
    ABS = ToAbs
    REAL = ToReal
    IMAG = ToImag

class ShotOrder(Enum):
    """Shot order allowed values.

    Generally, there are two possible modes in which a backend measures m
    circuits with n shots:
        - In the "circuit_first" mode, the backend subsequently first measures
          all m circuits and then repeats this n times.
        - In the "shot_first" mode, the backend first measures the 1st circuit
          n times, then the 2nd circuit n times, and it proceeds with the remaining
          circuits in the same way until it measures the m-th circuit n times.

    The current default mode of IBM Quantum devices is "circuit_first".
    """

    # pylint: disable=invalid-name
    circuit_first = "c"
    shot_first = "s"

class RestlessNode(DataAction, ABC):
    """An abstract node for restless data processing nodes.

    In restless measurements, the qubit is not reset after each measurement. Instead, the
    outcome of the previous quantum non-demolition measurement is the initial state for the
    current circuit. Restless measurements therefore require special data processing nodes
    that are implemented as sub-classes of `RestlessNode`. Restless experiments provide a
    fast alternative for several calibration and characterization tasks, for details
    see https://arxiv.org/pdf/2202.06981.pdf.

    This node takes as input an array of arrays (2d array) where the sub-arrays are
    the memories of each measured circuit. The sub-arrays therefore have a length
    given by the number of shots. This data is reordered into a one dimensional array where
    the element at index j was the jth measured shot. This node assumes by default that
    a list of circuits :code:`[circ_1, cric_2, ..., circ_m]` is measured :code:`n_shots`
    times according to the following order:

    .. parsed-literal::

        [
            circuit 1 - shot 1,
            circuit 2 - shot 1,
            ...
            circuit m - shot 1,
            circuit 1 - shot 2,
            circuit 2 - shot 2,
            ...
            circuit m - shot 2,
            circuit 1 - shot 3,
            ...
            circuit m - shot n,
        ]

    Once the shots have been ordered in this fashion the data can be post-processed.
    """

    def __init__(
        self, validate: bool = True, memory_allocation: ShotOrder = ShotOrder.circuit_first
    ):
        """Initialize a restless node.

        Args:
            validate: If set to True the node will validate its input.
            memory_allocation: If set to "c" the node assumes that the backend
                subsequently first measures all circuits and then repeats this
                n times, where n is the total number of shots. The default value
                is "c". If set to "s" it is assumed that the backend subsequently
                measures each circuit n times.
        """
        super().__init__(validate)
        self._n_shots = None
        self._n_circuits = None
        self._memory_allocation = memory_allocation

    def _format_data(self, data: np.ndarray) -> np.ndarray:
        """Convert the data to an array.

        This node will also set all the attributes needed to process the data such as
        the number of shots and the number of circuits.

        Args:
            data: An array representing the memory.

        Returns:
            The data that has been processed.

        Raises:
            DataProcessorError: If the datum has the wrong shape.
        """

        self._n_shots = len(data[0])
        self._n_circuits = len(data)

        if self._validate:
            if data.shape != (self._n_circuits, self._n_shots):
                raise DataProcessorError(
                    f"The datum given to {self.__class__.__name__} does not convert "
                    "of an array with dimension (number of circuit, number of shots)."
                )

        return data

    def _reorder(self, unordered_data: np.ndarray) -> np.ndarray:
        """Reorder the measured data according to the measurement sequence.

        Here, by default, it is assumed that the inner loop of the measurement
        is done over the circuits and the outer loop is done over the shots.
        The returned data is a one-dimensional array of time-ordered shots.
        """
        if unordered_data is None:
            return unordered_data

        if self._memory_allocation == ShotOrder.circuit_first:
            return unordered_data.T.flatten()
        else:
            return unordered_data.flatten()

class RestlessToCounts(RestlessNode):
    """Post-process restless data and convert restless memory to counts.

    This node first orders the measured restless data according to the measurement
    sequence and then compares each bit in a shot with its value in the previous shot.
    If they are the same then the bit corresponds to a 0, i.e. no state change, and if
    they are different then the bit corresponds to a 1, i.e. there was a state change.
    """

    def __init__(self, num_qubits: int, validate: bool = True):
        """
        Args:
            num_qubits: The number of qubits which is needed to construct the header needed
                by :code:`qiskit.result.postprocess.format_counts_memory` to convert the memory
                into a bit-string of counts.
            validate: If set to False the DataAction will not validate its input.
        """
        super().__init__(validate)
        self._num_qubits = num_qubits

    def _process(self, data: np.ndarray) -> np.ndarray:
        """Reorder the shots and assign values to them based on the previous outcome.

        Args:
            data: An array representing the memory.

        Returns:
            A counts dictionary processed according to the restless methodology.
        """

        # Step 1. Reorder the data.
        memory = self._reorder(data)

        # Step 2. Do the restless classification into counts.
        counts = [defaultdict(int) for _ in range(self._n_circuits)]
        prev_shot = "0" * self._num_qubits
        header = {"memory_slots": self._num_qubits}

        for idx, shot in enumerate(memory):
            shot = format_counts_memory(shot, header)

            restless_adjusted_shot = RestlessToCounts._restless_classify(shot, prev_shot)

            circuit_idx = idx % self._n_circuits

            counts[circuit_idx][restless_adjusted_shot] += 1

            prev_shot = shot

        return np.array([dict(counts_dict) for counts_dict in counts])

    @staticmethod
    def _restless_classify(shot: str, prev_shot: str) -> str:
        """Adjust the measured shot based on the previous shot.

        Each bit in shot is compared to its value in the previous shot. If both are equal
        the restless adjusted bit is 0 (no state change) otherwise it is 1 (the
        qubit changed state). This corresponds to taking the exclusive OR operation
        between each bit and its previous outcome.

        Args:
            shot: A measured shot as a binary string, e.g. "0110100".
            prev_shot: The shot that was measured in the previous circuit.

        Returns:
            The restless adjusted string computed by comparing the shot with the previous shot.
        """
        restless_adjusted_bits = []

        for idx, bit in enumerate(shot):
            restless_adjusted_bits.append("0" if bit == prev_shot[idx] else "1")

        return "".join(restless_adjusted_bits)


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

datos01 = np.loadtxt("euler.dat").T
datos1 = np.loadtxt("rk4.dat").T
datos001 = np.loadtxt("leapfrog.dat").T

plt.figure(figsize=(10,10))

plt.subplot(3,3,1)
plt.plot(datos01[0],datos01[1])
plt.xlim(0,15)
plt.ylim(-1.5,1.5)
plt.xlabel("Time")
plt.ylabel("Position")
plt.title("Euler")

plt.subplot(3,3,2)
plt.plot(datos1[0],datos1[1])
plt.xlim(0,15)
plt.ylim(-1.5,1.5)
plt.xlabel("Time")
plt.title("RK")

plt.subplot(3,3,3)
plt.plot(datos001[0],datos001[1])
plt.xlim(0,15)
plt.ylim(-1.5,1.5)
plt.xlabel("Time")
plt.title("Leapfrog")

plt.subplot(3,3,4)
plt.plot(datos01[0],datos01[2])
plt.xlim(0,15)
plt.ylim(-1.5,1.5)
plt.xlabel("Time")
plt.ylabel("Velocity")

plt.subplot(3,3,5)
plt.plot(datos1[0],datos1[2])
plt.xlim(0,15)
plt.ylim(-1.5,1.5)
plt.xlabel("Time")

plt.subplot(3,3,6)
plt.plot(datos001[0],datos001[2])
plt.xlim(0,15)
plt.ylim(-1.5,1.5)
plt.xlabel("Time")

plt.subplot(3,3,7)
plt.plot(datos01[1],datos01[2])
plt.xlim(-2,2)
plt.ylim(-2,2)
plt.xlabel("Position")

plt.subplot(3,3,8)
plt.plot(datos1[1],datos1[2])
plt.xlim(-2,2)
plt.ylim(-2,2)
plt.xlabel("Position")

plt.subplot(3,3,9)
plt.plot(datos001[1],datos001[2])
plt.xlim(-2,2)
plt.ylim(-2,2)
plt.xlabel("Position")

plt.savefig("resultados.png")

ext: py

reason: import numpy


text:
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Benchmarks for `tf.data.Dataset.list_files()`."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from os import path
from os import makedirs
import shutil
import time
import tempfile

import numpy as np

from tensorflow.python.client import session
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.platform import test

class ListFilesBenchmark(test.Benchmark):
  """Benchmarks for `tf.data.Dataset.list_files()`."""

  def benchmark_nested_directories(self):
    tmp_dir = tempfile.mkdtemp()
    width = 1024
    depth = 16
    for i in range(width):
      for j in range(depth):
        new_base = path.join(tmp_dir, str(i),
                             *[str(dir_name) for dir_name in range(j)])
        makedirs(new_base)
        child_files = ['a.py', 'b.pyc'] if j < depth - 1 else ['c.txt', 'd.log']
        for f in child_files:
          filename = path.join(new_base, f)
          open(filename, 'w').close()
    patterns = [
        path.join(tmp_dir, path.join(*['**'
                                       for _ in range(depth)]), suffix)
        for suffix in ['*.txt', '*.log']
    ]
    deltas = []
    iters = 3
    for _ in range(iters):
      with ops.Graph().as_default():
        dataset = dataset_ops.Dataset.list_files(patterns)
        options = dataset_ops.Options()
        options.experimental_optimization.apply_default_optimizations = False
        dataset = dataset.with_options(options)
        next_element = dataset.make_one_shot_iterator().get_next()
        with session.Session() as sess:
          sub_deltas = []
          while True:
            try:
              start = time.time()
              sess.run(next_element)
              end = time.time()
              sub_deltas.append(end - start)
            except errors.OutOfRangeError:
              break
          deltas.append(sub_deltas)
    median_deltas = np.median(deltas, axis=0)
    self.report_benchmark(
        iters=iters,
        wall_time=np.sum(median_deltas),
        extras={
            'read first file:':
                median_deltas[0],
            'read second file:':
                median_deltas[1],
            'avg time for reading %d more filenames:' %
            (len(median_deltas) - 2):
                np.average(median_deltas[2:])
        },
        name='nested_directory(%d*%d)' % (width, depth))
    shutil.rmtree(tmp_dir, ignore_errors=True)

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

max_stars_repo_name: uve/tensorflow

max_issues_repo_name: uve/tensorflow

max_forks_repo_name: uve/tensorflow


text:
# -*- coding: utf-8 -*-
"""
Created on Mon Nov  8 16:01:33 2021

@author: shangfr
"""

import numpy as np
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OrdinalEncoder, LabelEncoder, StandardScaler
from sklearn.ensemble import AdaBoostClassifier, AdaBoostRegressor
from sklearn.metrics import roc_curve, auc, classification_report, mean_squared_error, r2_score

from model.ml_plot import plot_roc_curve

def model_cls(model_dict):

    data = model_dict['data']
    target = model_dict['target']
    features = model_dict['features']
    positive = model_dict["positive"]

    X = data[features]
    y = data[target]
    y = y == positive

    enc = OrdinalEncoder(
        handle_unknown='use_encoded_value', unknown_value=np.nan)
    le = LabelEncoder()

    X = enc.fit_transform(X)
    y = le.fit_transform(y)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, random_state=0)

    pipe = Pipeline([('scaler', StandardScaler()),
                     ('clf', AdaBoostClassifier())])

    pipe.fit(X_train, y_train)

    # Classification metrics
    report = classification_report(y_test, pipe.predict(X_test))

    y_pred = pipe.predict_proba(X_test)[:, 1]

    fpr, tpr, thresholds = roc_curve(y_test, y_pred)

    auc_value = f'{auc(fpr, tpr):.4f}'

    fig = plot_roc_curve(fpr, tpr, thresholds, auc_value, positive)

    return fig, report

def model_regr(model_dict):

    data = model_dict['data']
    target = model_dict['target']
    features = model_dict['features']

    X = data[features]
    y = data[target]
    enc = OrdinalEncoder(
        handle_unknown='use_encoded_value', unknown_value=np.nan)

    X = enc.fit_transform(X)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, random_state=0)

    pipe = Pipeline([('scaler', StandardScaler()),
                     ('regr', AdaBoostRegressor())])

    pipe.fit(X_train, y_train)
    # Regression metrics
    y_pred = pipe.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    result = {"mean_squared_error": mse, "r2_score": r2}
    return result

max_stars_repo_name: shangfr/Data-Follower

max_issues_repo_name: shangfr/Data-Follower

max_forks_repo_name: shangfr/Data-Follower


text:
import pandas as pd
import sys
#import getopt
from time import time
import json
from pyspark.sql import DataFrame, Column
from pyspark.sql.types import StructType, StructField, StringType,ArrayType,DataType
from pyspark.sql.functions import col,udf,struct,collect_list, lower, count,lit
from pyspark import SparkContext
from pyspark.sql import SparkSession
import csv
from datetime import date, timedelta
import numpy as np
from typing import Callable

# def spark_context(appname='cms', yarn=None, verbose=False, python_files=[]):
#         # define spark context, it's main object which allow
#         # to communicate with spark
#     if  python_files:
#         return SparkContext(appName=appname, pyFiles=python_files)
#     else:
#         return SparkContext(appName=appname)

# def spark_session(appName="log-parser"):
#     """
#     Function to create new spark session
#     """
#     sc = SparkContext(appName="log-parser")
#     return SparkSession.builder.config(conf=sc._conf).getOrCreate()

class py_or_udf:
    "Introducing — `py_or_udf` — a decorator that allows a method to act as either a regular python method or a pyspark UDF"
    def __init__(self, returnType : DataType=StringType()):
        self.spark_udf_type = returnType

    def __call__(self, func : Callable):
        def wrapped_func(*args, **kwargs):
            if any([isinstance(arg, Column) for arg in args]) or \
                any([isinstance(vv, Column) for vv in kwargs.values()]):
                return udf(func, self.spark_udf_type)(*args, **kwargs)

            else:
                return func(*args, **kwargs)

        return wrapped_func

@py_or_udf(returnType=StringType())
def clean_message(message,clean_short=False):

    import re
    message = re.sub(r'[^!"()*+,-\\\/:;<=>?@[\]^_`{|}\w\s]', '', message)  # removes any foreign character
    message=message.lower()
    message = re.sub(r'(\w+:)*\/\S*', ' ', message)  # removes filepath/url
    message=re.sub(r'/\((.*?)\)/g',' ',message)#removes strings in ()
    message = re.sub(r'\S+@\S+', ' ', message) #removes strings with @
    message = re.sub(r'(\w+-)*\w+((\.|:)\w+)*(\.|:)\w+', ' ', message) #removes no-textual strings with . : - inside
    message = re.sub(r'(\w+\d+\w*)|(\d+\w+)', ' ', message)#removes alpha-numerical string
    message = re.sub(r'\w+_\S+', ' ', message)#removes strings with underscore
    message = re.sub(r'\.\w+', ' ', message)#removes strings with .
    message = re.sub(r'-\S+-', ' ', message)#removes strings between -
    message = re.sub(r'([^:]+:\s(?=[^:]+:[^:]))', ' ', message) #nucleus selection
    message = re.sub(r'\S+\s=\s\S+', ' ', message)#removes patterns with string equality
    message = re.sub(r'\S+=(\w+)*', ' ', message)
    message = re.sub(r'\[\w+\]', ' ', message)# removes [string]
    message = re.sub(r'(\d+)', ' ', message)#removes digits
    message = re.sub(r'[^\w\s]', ' ', message) # removes punctuation
#         message = re.sub(r'\\', ' ', message)#removes \
    if(clean_short):
        message = re.sub(r'\s\w{1,2}\b(?<!\bno)', ' ', message) #removes strings up to 2 characters long except for no
    else:
        message = re.sub(r'\s[a-zA-Z]{1}(\s|\Z)', ' ', message) #removes one char string
    message = re.sub(r' +', r' ', message)# removes addictional whitespace
    message = re.sub(r'\A ', '', message)#removes whitespace at the beginning
    message = re.sub(r' \Z', '', message)#removes whitespace at the end

    return message

def tokenize(df_cleaned):
    """
    This function removes tokens appearing once and stopwords
    """
    from pyspark.ml.feature import Tokenizer
    from pyspark.sql.types import IntegerType
    tokenizer = Tokenizer(inputCol="error_message", outputCol="tokenized_message")
    countTokens = udf(lambda words: len(words), IntegerType())
    tokenized = tokenizer.transform(df_cleaned)
    tokenized=tokenized.select("error_message", "tokenized_message").withColumn("tokens", countTokens(col("tokenized_message")))
    return tokenized

def clean_tokenized(df_tokenized,stop_one=True):
    """
    This function removes stopwords (optional) and words occuring once
    """
    from pyspark.ml.feature import StopWordsRemover
    import pyspark.sql.functions as F

    #stopword list
    stopwords = ['the','with','a','an','but','of','on','to','all','has','have','been','for','in','it','its','itself',
                'this','that','those','these','is','are','were','was','be','being','having','had','does','did','doing',
                'and','if','about','again','then','so','too','by','error']

   #removes words occurring once
    if(stop_one):
         count_df = df_tokenized.withColumn('word', F.explode(F.col('tokenized_message')))\
         .groupBy('word')\
         .count()\
         .sort('count', ascending=False)
         stoplist=count_df.where(F.col('count')==1).select('word').rdd.flatMap(lambda x: x).collect()
         stopwords=stoplist+stopwords

    remover = StopWordsRemover(inputCol="tokenized_message",outputCol="cleaned_tok_message", stopWords=stopwords)
    df_tokenized=remover.transform(df_tokenized)

    return df_tokenized

class preprocMex(object):

    def __init__(self,spark,hpath,clean_short,stop_one):

        self.hpath=hpath
        self.spark=spark
        self.clean_short=clean_short
        self.stop_one=stop_one

    def process(self,verbose=False):

        """
        Parse fts records from HDFS
        """

        # create new spark DataFrame
        schema = StructType([StructField('data', StructType([StructField('t__error_message', StringType(), nullable=True)]))])
        df=self.spark.read.json(self.hpath, schema)
        df=df.select(col('data.t__error_message').alias('error_message')).where('error_message <> ""')
        df.cache()
        bf_n=df.count()
        #print('before cleaning %i message'% bf_n)
        #print('...cleaning message')
        df=df.withColumn('error_message', clean_message(col('error_message'),lit(self.clean_short))).dropDuplicates()
        af_n=df.count()
        df=tokenize(df)
        df=clean_tokenized(df,stop_one=self.stop_one)
        #print('after cleaning %i different message'% af_n)
        #df.show()
        return df,bf_n,af_n

class MyCorpus(object):

    """An interator that yields sentences (lists of str)."""

    def __init__(self,inputDf):
        self.inputDf=inputDf
        self.list_err=self.inputDf.select(collect_list("cleaned_tok_message")).collect()[0][0]

    def __iter__(self):

        for line in self.list_err:
            yield line

max_stars_repo_name: micolocco/ClusterLog

max_issues_repo_name: micolocco/ClusterLog

max_forks_repo_name: micolocco/ClusterLog


text:
"""
Helper functions used by modules in `~halotools.mock_observables.pairwise_velocities`
to process arguments and raise appropriate exceptions.
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import numpy as np
from warnings import warn
from astropy.utils.misc import NumpyRNGContext

from ..mock_observables_helpers import (enforce_sample_has_correct_shape,
    get_period, get_num_threads)

from ...custom_exceptions import HalotoolsError
from ...utils.array_utils import array_is_monotonic

__all__ = ['_pairwise_velocity_stats_process_args', '_process_radial_bins', '_process_rp_bins']
__author__ = ['Duncan Campbell']

def _pairwise_velocity_stats_process_args(sample1, velocities1, sample2, velocities2,
        period, do_auto, do_cross, num_threads, approx_cell1_size, approx_cell2_size, seed):
    """
    Private method to do bounds-checking on the arguments passed to
    `~halotools.mock_observables.pairwise_velocity_stats`.
    """

    sample1 = enforce_sample_has_correct_shape(sample1)
    velocities1 = np.atleast_1d(velocities1).astype('f4')

    if sample2 is not None:
        sample2 = np.atleast_1d(sample2)
        if velocities2 is None:
            msg = ("\n If `sample2` is passed as an argument, \n"
                   "`velocities2` must also be specified.")
            raise ValueError(msg)
        else:
            velocities2 = np.atleast_1d(velocities2)
        if np.all(sample1 == sample2):
            _sample1_is_sample2 = True
            msg = ("\n Warning: `sample1` and `sample2` are exactly the same, \n"
                   "only the auto-function will be returned.\n")
            warn(msg)
            do_cross = False
        else:
            _sample1_is_sample2 = False
    else:
        sample2 = sample1
        velocities2 = velocities1
        _sample1_is_sample2 = True

    period, PBCs = get_period(period)

    if (type(do_auto) is not bool) | (type(do_cross) is not bool):
        msg = ('\n `do_auto` and `do_cross` keywords must be of type boolean.')
        raise HalotoolsError(msg)
    else:
        if (do_auto is False) & (do_cross is False):
            msg = ("Both ``do_auto`` and ``do_cross`` have been set to False")
            raise ValueError(msg)

    num_threads = get_num_threads(num_threads)

    return sample1, velocities1, sample2, velocities2, period, do_auto,\
        do_cross, num_threads, _sample1_is_sample2, PBCs

def _process_radial_bins(rbins, period, PBCs):
    """
    process radial bin parameter
    """

    # check the radial bins parameter
    rbins = np.atleast_1d(rbins)
    rmax = np.max(rbins)
    try:
        assert rbins.ndim == 1
        assert len(rbins) > 1
        if len(rbins) > 2:
            assert array_is_monotonic(rbins, strict=True) == 1
    except AssertionError:
        msg = ("\n Input `rbins` must be a monotonically increasing \n"
               "1-D array with at least two entries.")
        raise ValueError(msg)

    # check for input parameter consistency
    if PBCs:
        if (rmax >= np.min(period)/3.0):
            msg = ("\n The maximum length over which you search for pairs \n"
                   "of points cannot be larger than Lbox/3 in any dimension. \n"
                   "If you need to count pairs on these length scales, \n"
                   "you should use a larger simulation. \n")
            raise ValueError(msg)

    return rbins

def _process_rp_bins(rp_bins, pi_max, period, PBCs):
    """
    process projected radial bin and pi_max parameters
    """

    # process projected radial bins
    rp_bins = np.atleast_1d(rp_bins)
    rp_max = np.max(rp_bins)
    try:
        assert rp_bins.ndim == 1
        assert len(rp_bins) > 1
        if len(rp_bins) > 2:
            assert array_is_monotonic(rp_bins, strict=True) == 1
    except AssertionError:
        msg = ("\n Input `rp_bins` must be a monotonically increasing \n"
               "1-D array with at least two entries.")
        raise ValueError(msg)

    pi_max = float(pi_max)

    if PBCs:
        if (rp_max >= np.min(period[0:2])/3.0):
            msg = ("\n The maximum length over which you search for pairs of points \n"
                   "cannot be larger than Lbox/3 in any dimension. \n"
                   "If you need to count pairs on these length scales, \n"
                   "you should use a larger simulation.\n")
            raise ValueError(msg)
        if (pi_max >= np.min(period[2])/3.0):
            msg = ("\n The maximum length over which you search for pairs of points \n"
                   "cannot be larger than Lbox/3 in any dimension. \n"
                   "The input ``pi_max`` = {0}"
                   "If you need to count pairs on these length scales, \n"
                   "you should use a larger simulation.\n".format(pi_max))
            raise ValueError(msg)

    return rp_bins, pi_max

max_stars_repo_name: pllim/halotools

max_issues_repo_name: pllim/halotools

max_forks_repo_name: pllim/halotools

reason: import numpy,from astropy


text:
# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s

import numpy as np
import os
import sys

_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
sys.path.append(_SCRIPT_PATH)
from tools import mlir_pytaco_api as pt

compressed = pt.compressed

i, j = pt.get_index_vars(2)
A = pt.tensor([2, 3])
S = pt.tensor(3) # S is a scalar tensor.
B = pt.tensor([2, 3], compressed)
A.insert([0, 1], 10)
A.insert([1, 2], 40)

# Use [0] to index the scalar tensor.
B[i, j] = A[i, j] * S[0]

indices, values = B.get_coordinates_and_values()
passed = np.array_equal(indices, [[0, 1], [1, 2]])
passed += np.array_equal(values, [30.0, 120.0])

# Sum all the values in A.
S[0] = A[i, j]
passed += (S.get_scalar_value() == 50.0)

indices, values = S.get_coordinates_and_values()
passed += (len(indices)==0)
passed += (values == 50.0)

# CHECK: Number of passed: 5
print("Number of passed:", passed)

max_stars_repo_name: lxbndr/llvm-project

max_issues_repo_name: lxbndr/llvm-project

max_forks_repo_name: lxbndr/llvm-project


text:
# -*- coding: utf-8 -*-
# Copyright 2018 the HERA Project
# Licensed under the MIT License

import unittest
from hera_cal import datacontainer, io
import numpy as np
from hera_cal.data import DATA_PATH
import os
from hera_cal import abscal

class TestDataContainer(unittest.TestCase):

    def setUp(self):
        self.antpairs = [(1, 2), (2, 3), (3, 4), (1, 3), (2, 4)]  # not (1,4)
        self.pols = ['xx', 'yy']
        self.blpol = {}
        for bl in self.antpairs:
            self.blpol[bl] = {}
            for pol in self.pols:
                self.blpol[bl][pol] = 1j
        self.polbl = {}
        for pol in self.pols:
            self.polbl[pol] = {}
            for bl in self.antpairs:
                self.polbl[pol][bl] = 1j
        self.both = {}
        for pol in self.pols:
            for bl in self.antpairs:
                self.both[bl + (pol,)] = 1j

    def test_init(self):
        dc = datacontainer.DataContainer(self.blpol)
        for k in dc._data.keys():
            self.assertEqual(len(k), 3)
        self.assertEqual(set(self.antpairs), dc._antpairs)
        self.assertEqual(set(self.pols), dc._pols)
        dc = datacontainer.DataContainer(self.polbl)
        for k in dc._data.keys():
            self.assertEqual(len(k), 3)
        self.assertEqual(set(self.antpairs), dc._antpairs)
        self.assertEqual(set(self.pols), dc._pols)
        dc = datacontainer.DataContainer(self.both)
        for k in dc._data.keys():
            self.assertEqual(len(k), 3)
        self.assertEqual(set(self.antpairs), dc._antpairs)
        self.assertEqual(set(self.pols), dc._pols)
        self.assertIsNone(dc.antpos)
        self.assertIsNone(dc.freqs)
        self.assertIsNone(dc.times)
        self.assertIsNone(dc.lsts)
        self.assertIsNone(dc.times_by_bl)
        self.assertIsNone(dc.lsts_by_bl)
        self.assertRaises(
            AssertionError, datacontainer.DataContainer, {(1, 2, 3, 4): 2})

    def test_antpairs(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertEqual(set(self.antpairs), dc.antpairs())
        self.assertEqual(set(self.antpairs), dc.antpairs('xx'))
        self.assertEqual(set(self.antpairs), dc.antpairs('yy'))
        dc = datacontainer.DataContainer(self.polbl)
        self.assertEqual(set(self.antpairs), dc.antpairs())
        self.assertEqual(set(self.antpairs), dc.antpairs('xx'))
        self.assertEqual(set(self.antpairs), dc.antpairs('yy'))
        dc = datacontainer.DataContainer(self.both)
        self.assertEqual(set(self.antpairs), dc.antpairs())
        self.assertEqual(set(self.antpairs), dc.antpairs('xx'))
        self.assertEqual(set(self.antpairs), dc.antpairs('yy'))

    def test_bls(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertEqual(set(dc.keys()), dc.bls())
        dc = datacontainer.DataContainer(self.polbl)
        self.assertEqual(set(dc.keys()), dc.bls())
        dc = datacontainer.DataContainer(self.both)
        self.assertEqual(set(dc.keys()), dc.bls())

    def test_pols(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertEqual(set(self.pols), dc.pols())
        self.assertEqual(set(self.pols), dc.pols((1, 2)))
        dc = datacontainer.DataContainer(self.polbl)
        self.assertEqual(set(self.pols), dc.pols())
        self.assertEqual(set(self.pols), dc.pols((1, 2)))
        dc = datacontainer.DataContainer(self.both)
        self.assertEqual(set(self.pols), dc.pols())
        self.assertEqual(set(self.pols), dc.pols((1, 2)))

    def test_keys(self):
        dc = datacontainer.DataContainer(self.blpol)
        keys = dc.keys()
        self.assertEqual(len(keys), len(self.pols) * len(self.antpairs))
        dc = datacontainer.DataContainer(self.polbl)
        keys = dc.keys()
        self.assertEqual(len(keys), len(self.pols) * len(self.antpairs))
        dc = datacontainer.DataContainer(self.both)
        keys = dc.keys()
        self.assertEqual(len(keys), len(self.pols) * len(self.antpairs))

    def test_values(self):
        dc = datacontainer.DataContainer(self.blpol)
        values = dc.values()
        self.assertEqual(len(values), len(self.pols) * len(self.antpairs))
        self.assertEqual(values[0], 1j)
        dc = datacontainer.DataContainer(self.polbl)
        values = dc.values()
        self.assertEqual(len(values), len(self.pols) * len(self.antpairs))
        self.assertEqual(values[0], 1j)
        dc = datacontainer.DataContainer(self.both)
        values = dc.values()
        self.assertEqual(len(values), len(self.pols) * len(self.antpairs))
        self.assertEqual(values[0], 1j)

    def test_items(self):
        dc = datacontainer.DataContainer(self.blpol)
        items = dc.items()
        self.assertEqual(len(items), len(self.pols) * len(self.antpairs))
        self.assertTrue(items[0][0][0:2] in self.antpairs)
        self.assertTrue(items[0][0][2] in self.pols)
        self.assertEqual(items[0][1], 1j)
        dc = datacontainer.DataContainer(self.polbl)
        items = dc.items()
        self.assertTrue(items[0][0][0:2] in self.antpairs)
        self.assertTrue(items[0][0][2] in self.pols)
        self.assertEqual(items[0][1], 1j)
        dc = datacontainer.DataContainer(self.both)
        items = dc.items()
        self.assertTrue(items[0][0][0:2] in self.antpairs)
        self.assertTrue(items[0][0][2] in self.pols)
        self.assertEqual(items[0][1], 1j)

    def test_len(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertEqual(len(dc), 10)
        dc = datacontainer.DataContainer(self.polbl)
        self.assertEqual(len(dc), 10)
        dc = datacontainer.DataContainer(self.both)
        self.assertEqual(len(dc), 10)

    def test_del(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertTrue((1, 2, 'xx') in dc)
        self.assertTrue((1, 2, 'XX') in dc)
        del dc[(1, 2, 'xx')]
        self.assertFalse((1, 2, 'xx') in dc)
        self.assertTrue('xx' in dc.pols())
        self.assertTrue((1, 2) in dc.antpairs())
        del dc[(1, 2, 'yy')]
        self.assertFalse((1, 2) in dc.antpairs())
        del dc[(2, 3, 'XX')]
        self.assertFalse((2, 3, 'xx') in dc)
        self.assertTrue('xx' in dc.pols())
        self.assertTrue((2, 3) in dc.antpairs())

    def test_getitem(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertEqual(dc[(1, 2, 'xx')], 1j)
        self.assertEqual(dc[(2, 1, 'xx')], -1j)
        self.assertEqual(dc[(1, 2)], {'xx': 1j, 'yy': 1j})
        self.assertEqual(set(dc['xx'].keys()), set(self.antpairs))
        self.assertEqual(dc[(1, 2, 'xx')], dc.get_data((1, 2, 'xx')))
        self.assertEqual(dc[(1, 2, 'xx')], dc.get_data(1, 2, 'xx'))
        dc = datacontainer.DataContainer(self.polbl)
        self.assertEqual(dc[(1, 2, 'xx')], 1j)
        self.assertEqual(dc[(2, 1, 'xx')], -1j)
        self.assertEqual(dc[(1, 2)], {'xx': 1j, 'yy': 1j})
        self.assertEqual(set(dc['xx'].keys()), set(self.antpairs))
        self.assertEqual(dc[(2, 1, 'xx')], dc.get_data((2, 1, 'xx')))
        self.assertEqual(dc[(2, 1, 'xx')], dc.get_data(2, 1, 'xx'))
        dc = datacontainer.DataContainer(self.both)
        self.assertEqual(dc[(1, 2, 'xx')], 1j)
        self.assertEqual(dc[(2, 1, 'xx')], -1j)
        self.assertEqual(dc[(1, 2)], {'xx': 1j, 'yy': 1j})
        self.assertEqual(set(dc['xx'].keys()), set(self.antpairs))
        self.assertEqual(dc[(1, 2)], dc.get_data((1, 2)))
        self.assertEqual(dc[(1, 2)], dc.get_data(1, 2))
        self.assertEqual(dc[(1, 2, 'XX')], 1j)
        self.assertEqual(dc[(2, 1, 'XX')], -1j)
        self.assertEqual(dc[(2, 1, 'XX')], dc.get_data(2, 1, 'XX'))
        self.assertEqual(dc[(2, 1, 'XX')], dc.get_data(2, 1, 'xx'))

    def test_has_key(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertTrue((2, 3, 'yy') in dc)
        self.assertTrue(dc.has_key((2, 3), 'yy'))
        self.assertTrue(dc.has_key((3, 2), 'yy'))
        self.assertFalse('xy' in dc)
        self.assertFalse((5, 6) in dc)
        self.assertFalse((1, 2, 'xy') in dc)
        dc = datacontainer.DataContainer(self.polbl)
        self.assertTrue((2, 3, 'yy') in dc)
        self.assertTrue(dc.has_key((2, 3), 'yy'))
        self.assertTrue(dc.has_key((3, 2), 'yy'))
        self.assertFalse('xy' in dc)
        self.assertFalse((5, 6) in dc)
        self.assertFalse((1, 2, 'xy') in dc)
        dc = datacontainer.DataContainer(self.both)
        self.assertTrue((2, 3, 'yy') in dc)
        self.assertTrue(dc.has_key((2, 3), 'yy'))
        self.assertTrue(dc.has_key((3, 2), 'yy'))
        self.assertFalse('xy' in dc)
        self.assertFalse((5, 6) in dc)
        self.assertFalse((1, 2, 'xy') in dc)

        self.assertTrue((2, 3, 'YY') in dc)
        self.assertTrue(dc.has_key((2, 3), 'YY'))
        self.assertTrue(dc.has_key((3, 2), 'YY'))
        self.assertFalse('XY' in dc)
        self.assertFalse((5, 6) in dc)
        self.assertFalse((1, 2, 'XY') in dc)

        # assert switch bl
        dc[(1, 2, 'xy')] = 1j
        self.assertTrue((2, 1, 'yx') in dc)

    def test_has_antpair(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertTrue(dc.has_antpair((2, 3)))
        self.assertTrue(dc.has_antpair((3, 2)))
        self.assertFalse(dc.has_antpair((0, 3)))
        dc = datacontainer.DataContainer(self.polbl)
        self.assertTrue(dc.has_antpair((2, 3)))
        self.assertTrue(dc.has_antpair((3, 2)))
        self.assertFalse(dc.has_antpair((0, 3)))
        dc = datacontainer.DataContainer(self.both)
        self.assertTrue(dc.has_antpair((2, 3)))
        self.assertTrue(dc.has_antpair((3, 2)))
        self.assertFalse(dc.has_antpair((0, 3)))

    def test_has_pol(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertTrue(dc.has_pol('xx'))
        self.assertTrue(dc.has_pol('XX'))
        self.assertFalse(dc.has_pol('xy'))
        self.assertFalse(dc.has_pol('XY'))
        dc = datacontainer.DataContainer(self.polbl)
        self.assertTrue(dc.has_pol('xx'))
        self.assertTrue(dc.has_pol('XX'))
        self.assertFalse(dc.has_pol('xy'))
        self.assertFalse(dc.has_pol('XY'))
        dc = datacontainer.DataContainer(self.both)
        self.assertTrue(dc.has_pol('xx'))
        self.assertTrue(dc.has_pol('XX'))
        self.assertFalse(dc.has_pol('xy'))
        self.assertFalse(dc.has_pol('XY'))

    def test_get(self):
        dc = datacontainer.DataContainer(self.blpol)
        self.assertEqual(dc.get((1, 2), 'yy'), 1j)
        self.assertEqual(dc.get((2, 1), 'yy'), -1j)
        dc = datacontainer.DataContainer(self.polbl)
        self.assertEqual(dc.get((1, 2), 'yy'), 1j)
        self.assertEqual(dc.get((2, 1), 'yy'), -1j)
        dc = datacontainer.DataContainer(self.both)
        self.assertEqual(dc.get((1, 2), 'yy'), 1j)
        self.assertEqual(dc.get((2, 1), 'yy'), -1j)
        self.assertEqual(dc.get((1, 2), 'YY'), 1j)
        self.assertEqual(dc.get((2, 1), 'YY'), -1j)

    def test_setter(self):
        dc = datacontainer.DataContainer(self.blpol)
        # test basic setting
        dc[(100, 101, 'xy')] = np.arange(100) + np.arange(100) * 1j
        self.assertEqual(dc[(100, 101, 'xy')].shape, (100,))
        self.assertEqual(dc[(100, 101, 'xy')].dtype, np.complex)
        self.assertAlmostEqual(dc[(100, 101, 'xy')][1], (1 + 1j))
        self.assertAlmostEqual(dc[(101, 100, 'yx')][1], (1 - 1j))
        self.assertEqual(len(dc.keys()), 11)
        self.assertEqual((100, 101) in dc._antpairs, True)
        self.assertEqual('xy' in dc._pols, True)
        # test error
        self.assertRaises(ValueError, dc.__setitem__, *((100, 101), 100j))

    def test_adder(self):
        test_file = os.path.join(DATA_PATH, "zen.2458043.12552.xx.HH.uvORA")
        d, f = io.load_vis(test_file, pop_autos=True)
        d2 = d + d
        self.assertAlmostEqual(d2[(24, 25, 'xx')][30, 30], d[(24, 25, 'xx')][30, 30] * 2)
        # test exception
        d2, f2 = io.load_vis(test_file, pop_autos=True)
        d2[d2.keys()[0]] = d2[d2.keys()[0]][:, :10]
        self.assertRaises(ValueError, d.__add__, d2)
        d2[d2.keys()[0]] = d2[d2.keys()[0]][:10, :]
        self.assertRaises(ValueError, d.__add__, d2)

    def test_mul(self):
        test_file = os.path.join(DATA_PATH, "zen.2458043.12552.xx.HH.uvORA")
        d, f = io.load_vis(test_file, pop_autos=True)
        f[(24, 25, 'xx')][:, 0] = False
        f2 = f * f
        self.assertFalse(f2[(24, 25, 'xx')][0, 0])
        # test exception
        d2, f2 = io.load_vis(test_file, pop_autos=True)
        d2[d2.keys()[0]] = d2[d2.keys()[0]][:, :10]
        self.assertRaises(ValueError, d.__mul__, d2)
        d2[d2.keys()[0]] = d2[d2.keys()[0]][:10, :]
        self.assertRaises(ValueError, d.__mul__, d2)

    def test_concatenate(self):
        test_file = os.path.join(DATA_PATH, "zen.2458043.12552.xx.HH.uvORA")
        d, f = io.load_vis(test_file, pop_autos=True)
        d2 = d.concatenate(d)
        self.assertEqual(d2[(24, 25, 'xx')].shape[0], d[(24, 25, 'xx')].shape[0] * 2)
        d2 = d.concatenate(d, axis=1)
        self.assertEqual(d2[(24, 25, 'xx')].shape[1], d[(24, 25, 'xx')].shape[1] * 2)
        d2 = d.concatenate([d, d], axis=0)
        self.assertEqual(d2[(24, 25, 'xx')].shape[0], d[(24, 25, 'xx')].shape[0] * 3)
        # test exceptions
        d2, f2 = io.load_vis(test_file, pop_autos=True)
        d2[d2.keys()[0]] = d2[d2.keys()[0]][:10, :10]
        self.assertRaises(ValueError, d.concatenate, d2, axis=0)
        self.assertRaises(ValueError, d.concatenate, d2, axis=1)

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


text:
#
#  Copyright 2020 The FLEX Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import multiprocessing
from functools import partial

from typing import Dict, Tuple, Union

import numpy as np

from .. import gmpy_math
from .encrypted_number import PaillierEncryptedNumber
from .fixedpoint_number import FixedPointNumber
from .keypair import PaillierPublicKey, PaillierPrivateKey
from .obfuscator import apply_obfuscation
from .raw_encrypt import raw_encrypt

class PaillierEncryptor(object):
    def __init__(self, pub_key: PaillierPublicKey):
        self.pub_key = pub_key

    def __raw_encrypt(self, plaintext: int, random_value: int = None) -> int:
        """
        Encode an int into a encrypted number.
        Args:
            plaintext: int or float, as input number.
            random_value: int, used to do obfuscate if given.

        Returns:
            int, encrypt number.
        """
        return raw_encrypt(plaintext, self.pub_key, random_value)

    def _encrypt(self, value: Union[int, float], precision: int = None, random_value: int = None) -> PaillierEncryptedNumber:
        """Encode and Paillier encrypt a real number value.

        Args:
            value: int or float, as input number.
            precision: int, the numerical precision
            random_value: int, used to do obfuscate if given.

        Returns:
            PaillierEncryptedNumber
        """
        encoding = FixedPointNumber.encode(
            value, self.pub_key.n, self.pub_key.max_int, precision)
        obfuscator = random_value or 1
        ciphertext = self.__raw_encrypt(
            encoding.encoding, random_value=obfuscator)
        encryptednumber = PaillierEncryptedNumber(
            self.pub_key, ciphertext, encoding.exponent)
        if random_value is None:
            encryptednumber.apply_obfuscation()

        return encryptednumber

    def _encrypt_numpy(self, values_numpy: np.ndarray, precision: int = None, random_value: int = None) -> np.ndarray:
        """
        Traverse and encrypt every number in input numpy.array.

        Args:
            values_numpy: numpy.array. Should be int or float type.
            precision: int, the numerical precision
            random_value: int, used to do obfuscate if given.

        Returns:
            Numpy ndarray with the same shape.
        """
        s = values_numpy.shape
        values_flatten = values_numpy.reshape(-1)
        # Use multi process to accelerate computaion.
        if len(values_flatten) < 64:
            encryptednumber_flatten = np.array(
                [self._encrypt(v, precision, random_value) for v in values_flatten])
        else:
            process_num = min(multiprocessing.cpu_count(), len(values_flatten))
            pool = multiprocessing.Pool(processes=process_num)
            partial_encrypt = partial(
                self._encrypt, precision=precision, random_value=random_value)
            encryptednumber_flatten = np.array(
                pool.map(partial_encrypt, values_flatten))
            pool.terminate()
        return encryptednumber_flatten.reshape(s)

    def encrypt(self, value: Union[np.ndarray, int, float], precision: int = None, random_value: int = None) -> Union[np.ndarray, PaillierEncryptedNumber]:
        """
        Proxy for encrypt _encrypt and _encrypt_numpy. Use proper method to encrypt input.

        Args:
            value: int, float or numpy.array[Should be int or float type].
            precision: int, the numerical precision.
            random_value: int, used to do obfuscate if given.

        Returns:
            PaillierEncryptedNumber or Numpy ndarray of PaillierEncryptedNumber with the same shape as input.
        """
        if isinstance(value, np.ndarray):
            return self._encrypt_numpy(value, precision, random_value)
        else:
            return self._encrypt(value, precision, random_value)


text:
"""
Main class for creating a NavBoxPlus object, which provides a UKF for state and parameter
estimation as well as a general framework to feed those estimates into an adaptive controller.

"""
from __future__ import division
import numpy as np; npl = np.linalg
import scipy.linalg as spl

class NavBoxPlus(object):
    """
    Use the predict and correct functions as measurements come in.
    Note that the predict function will also return the control you should apply.
    Call get_state_and_cov to, you know, get the current state estimate and covariance.

    """
    def __init__(self, x0, Cx0, g, f, hDict, n_r, n_wf, n_whDict,
                 xplus=np.add, xminus=np.subtract, zminusDict=None,
                 alpha=1E-3, beta=2, kappa=1E-6):
        """
                x0: Best guess at the initial full state, (n_x >= n_m)
               Cx0: Initial state estimate covariance matrix, (n_m by n_m)
                 g: Control generating function, u = g(r, rnext, x, Cx, dt)
                 f: Discrete-time state dynamic, xnext = f(xlast, u, wf, dt)
             hDict: Dictionary of sensor model functions, {'sensor_i': zi = hi(x, u, whi)}
               n_r: Number of physical (non-parameter) state values
              n_wf: Number of values in process noise array
          n_whDict: Dictionary like hDict but containing the lengths of the sensor noise arrays {'sensor_i': n_whi}
             xplus: Perturbs a state by a tangent n_m-vector v on the state manifold, x2 = xplus(x1, v)
            xminus: Returns the tangent n_m-vector v that perturbs one state to another, v = xminus(x2, x1)
        zminusDict: Dictionary of measurement boxminus functions (like xminus is for the state), None == np.subtract
             alpha: |
              beta: | Unscented transform parameters (Wikipedia's recommendations)
             kappa: |

        Dimensionalities of all other variables are deduced from these inputs.

        """
        # Initialize state and draw the distinction between physical states and parameters
        self.x = np.atleast_1d(x0).astype(np.float64)
        self.n_x = len(self.x)
        self.n_r = int(n_r)
        self.n_p = self.n_x - self.n_r
        assert self.n_p >= 0

        # Initialize state covariance matrix and infer state-space manifold dimensionality
        self.Cx = np.atleast_2d(Cx0).astype(np.float64)
        self.n_m = len(self.Cx)
        assert self.n_x >= self.n_m

        # Verify that state covariance is positive definite
        self.is_pdef = lambda M: np.all(npl.eigvals(M) > 0)
        assert self.is_pdef(self.Cx)

        # Set controller function and infer control dimensionality
        self.g = g
        _len1d = lambda a: len(np.atleast_1d(a))
        self.n_u = _len1d(self.g(np.random.sample(self.n_r), np.random.sample(self.n_r),
                                 np.random.sample(self.n_x), self.Cx, np.random.sample()))
        self.u = np.zeros(self.n_u)

        # Set and verify process model
        self.f = f
        self.n_wf = int(n_wf)
        assert _len1d(self.f(np.random.sample(self.n_x),
                             np.random.sample(self.n_u),
                             np.random.sample(self.n_wf),
                             np.random.sample())) == self.n_x

        # Store sensor dictionaries and key list
        self.hDict = hDict
        self.hKeys = self.hDict.keys()
        self.n_whDict = {}
        self.n_zDict = {}
        self.zminusDict = {}
        for key in self.hKeys:
            self.n_whDict[key] = int(n_whDict[key])
            self.n_zDict[key] = _len1d(self.hDict[key](np.random.sample(self.n_x),
                                                       np.random.sample(self.n_u),
                                                       np.random.sample(self.n_whDict[key])))
            try: self.zminusDict[key] = zminusDict[key]
            except: self.zminusDict[key] = np.subtract

        # Set and verify boxplus
        self.xplus = xplus
        assert _len1d(self.xplus(np.random.sample(self.n_x), np.random.sample(self.n_m))) == self.n_x

        # Set and verify boxminus
        self.xminus = xminus
        assert _len1d(self.xminus(np.random.sample(self.n_x), np.random.sample(self.n_x))) == self.n_m

        # Process and sensor models wrapped for augmented-state sigma points
        self.sf = lambda s, dt: self.f(s[:self.n_x], self.u, s[self.n_x:], dt)
        self.sh = lambda s, h: h(s[:self.n_x], self.u, s[self.n_x:])

        # Boxoperations wrapped for augmented-state sigma points
        self.sxplus = lambda s, sv: np.append(self.xplus(s[:self.n_x], sv[:self.n_m]),
                                              s[self.n_x:] + sv[self.n_m:])
        self.sxminus = lambda s2, s1: np.append(self.xminus(s2[:self.n_x], s1[:self.n_x]),
                                                s2[self.n_x:] - s1[self.n_x:])

        # Core unscented transform parameters
        self.alpha = np.float64(alpha)
        self.beta = np.float64(beta)
        self.kappa = np.float64(kappa)

        # Memoize derived unscented transform parameters
        self.utpDict = {}
        alpha2 = self.alpha**2
        bp1ma2 = self.beta + 1 - alpha2
        a2m1 = alpha2 - 1
        a2k = alpha2*self.kappa
        for key in ['_f'] + self.hKeys:
            if key == '_f': L = self.n_wf + self.n_m
            else: L = self.n_whDict[key] + self.n_m
            l = L*a2m1 + a2k
            Lpl = L + l
            ldLpl = l/Lpl
            r2Lpl = 1/(2*Lpl)
            lba = ldLpl + bp1ma2
            self.utpDict[key] = [Lpl, ldLpl, r2Lpl, lba]

        # Super sensor dictionary contains all sensor metadata
        self.hDict_full = {}
        for key in self.hKeys:
            self.hDict_full[key] = (self.hDict[key],
                                    self.zminusDict[key],
                                    self.n_zDict[key],
                                    self.n_whDict[key],
                                    self.utpDict[key])

    def predict(self, r, rnext, wf0, Cf, dt):
        """
            r: Desired reference state at time t
        rnext: Desired reference state at time t+dt
          wf0: Mean of the process noise
           Cf: Covariance of the process noise
           dt: Timestep to predict forward

        Progresses the state estimate one dt into the future.
        Returns the control effort u.

        """
        # Compute control action
        self.u = self.g(r, rnext, self.x, self.Cx, dt)

        # Store relevant parameters
        utp = self.utpDict['_f']

        # Compute sigma points, propagate through process, and store tangent-space representation
        M = spl.cholesky(utp[0]*spl.block_diag(self.Cx, Cf))
        s0 = np.append(self.x, wf0)
        fS = [self.sf(s0, dt)]
        fT_sum = np.zeros(self.n_m)
        for Vi in np.vstack((M, -M)):
            fS.append(self.sf(self.sxplus(s0, Vi), dt))
            fT_sum += self.xminus(fS[-1], fS[0])

        # Determine the mean of the propagated sigma points from the tangent vectors
        self.x = self.xplus(fS[0], utp[2]*fT_sum)

        # Determine the covariance from the tangent-space deviations from the mean
        fv0 = self.xminus(fS[0], self.x)
        fPi_sum = np.zeros((self.n_m, self.n_m))
        for fSi in fS[1:]:
            fv = self.xminus(fSi, self.x)
            fPi_sum += np.outer(fv, fv)
        self.Cx = utp[3]*np.outer(fv0, fv0) + utp[2]*fPi_sum

        # Over and out
        return np.copy(self.u)

    def correct(self, sensor, z, wh0, Ch):
        """
        sensor: String of the sensor key name
             z: Value of the measurement
           wh0: Mean of that sensor's noise
            Ch: Covariance of that sensor's noise

        Updates the state estimate with the given sensor measurement.

        """
        # Store relevant functions and parameters
        h, zminus, n_z, _, utp = self.hDict_full[sensor]

        # Compute sigma points and emulate their measurement error
        M = spl.cholesky(utp[0]*spl.block_diag(self.Cx, Ch))
        V = np.vstack((M, -M))
        s0 = np.append(self.x, wh0)
        hE = [zminus(z, self.sh(s0, h))]
        for i, Vi in enumerate(V):
            hE.append(zminus(z, self.sh(self.sxplus(s0, Vi), h)))
        hE = np.array(hE)

        # Determine the mean of the sigma measurement errors
        e0 = utp[1]*hE[0] + utp[2]*np.sum(hE[1:], axis=0)

        # Determine the covariance and cross-covariance
        hV = hE - e0
        hPi_sum = np.zeros((n_z, n_z))
        hPci_sum = np.zeros((self.n_m, n_z))
        for Vqi, hVi in zip(V[:, :self.n_m], hV[1:]):
            hPi_sum += np.outer(hVi, hVi)
            hPci_sum += np.outer(Vqi, hVi)
        Pzz = utp[3]*np.outer(hV[0], hV[0]) + utp[2]*hPi_sum
        Pxz = utp[2]*hPci_sum

        # Kalman update the state estimate and covariance
        K = Pxz.dot(npl.inv(Pzz))
        self.x = self.xplus(self.x, -K.dot(e0))
        self.Cx = self.Cx - K.dot(Pzz).dot(K.T)

    def get_state_and_cov(self):
        """
        Returns the current state estimate and its covariance matrix.

        """
        return np.copy(self.x), np.copy(self.Cx)

reason: import numpy,import scipy


text:
import numpy as np
from numpy import linalg
from algos.GLASSO.base import base
from utils.common import np_soft_threshold
from utils.GLASSO.glasso import objective_F_cholesky

class OBN(base):
    def __init__(self, T,  N, lam, inner_T, ls_iter, step_lim):
        super(OBN, self).__init__(T, N, lam)
        self.inner_T = inner_T
        self.ls_iter = ls_iter
        self.step_lim = step_lim
        self.save_name = "OBN_N{N}_T{T}_innerT{inner_T}_LsIter{ls_iter}_StepLim{step_lim}" \
            .format(N=self.N, T=self.T, inner_T=self.inner_T, ls_iter=self.ls_iter, step_lim=self.step_lim)

    def compute(self, S, A0, status_f, history, test_check_f):
        As = []
        status = []

        lam = np.float32(self.lam)
        init_step = 1.0

        if A0 is None:
            A_diag = self.lam * np.ones(self.N, dtype='float32')
            A_diag = A_diag + np.diag(S)
            A_diag = 1.0 / A_diag
            A = np.diag(A_diag)
            A_diag = None
        else:
            A = np.array(A0, dtype='float32')

        if history:
            As.append(A.copy())

        if status_f is not None: status.append(status_f(A, 0.0))

        for t in range(self.T):
            A_inv = np.linalg.inv(A)
            if test_check_f is not None:
                if test_check_f(A, S, self.lam, A_inv):
                    t -= 1
                    break

            if init_step != 0:
                sign_A = np.sign(A, dtype='float32')
                mask_A = np.abs(sign_A, dtype='float32').astype('int8')
                G = S - A_inv
                G_min = np_soft_threshold(G + lam*sign_A, lam*(1.0 - mask_A))
                sign_soft_G = np.sign(np_soft_threshold(G, lam), dtype='float32')
                mask_G = np.abs(sign_soft_G).astype('int8')
                mask = np.bitwise_or(mask_A, mask_G)
                Z = sign_A - np.bitwise_xor(mask, mask_A) * sign_soft_G
                sign_A = None
                sign_soft_G = None
                mask_G = None
                mask_A = None

                X = np.zeros(S.shape, dtype='float32')
                R = - mask * (G + lam*Z)
                G = None
                Q = R
                RR_old = np.sum(R * R, dtype='float32')
                epsilon = 1e-10
                inner_T = self.inner_T
                if inner_T < 0:
                    inner_T = 5 + t//(-inner_T)
                    if t % (-self.inner_T) == 0: init_step = 1.0
                for inner_t in range(np.minimum(S.shape[0],inner_T)):
                    if init_step == 0: break
                    if RR_old < epsilon: break
                    Y = mask * (A_inv @ Q @ A_inv)
                    alpha = RR_old / np.sum(Q * Y, dtype='float32')
                    X = X + alpha * Q
                    R = R - alpha * Y
                    RR = np.sum(R * R, dtype='float32')
                    beta = RR / RR_old
                    Q = R + beta * Q
                    RR_old = RR
                A_inv = None

                A, step = projected_linesearch_F(A, S, lam, G_min, Z, X, init=init_step, max_iter=self.ls_iter, step_lim=self.step_lim)
                if step == 0: init_step = 0

            if history:
                As.append(A.copy())

            if status_f is not None: status.append(status_f(A, step))

        if init_step == 0: t = np.inf
        return A, status, As, t+1

def projected_linesearch_F(A, S, lam, g_min, Z, Delta, init = 1, beta = 0.5, c = 0.1 , max_iter=10, step_lim=0):
    L = np.linalg.cholesky(A)
    init_F_val = objective_F_cholesky(A,S,lam,L)
    init_F_val_gA = init_F_val - c*np.sum(g_min*A, dtype='float32')
    init_F_val = None
    step = init
    for _ in range(max_iter):
        if step < step_lim: break
        try:
            A_next = A+step*Delta
            A_next[np.sign(A_next, dtype='float32') != Z] = 0
            A_next = A_next + np.transpose(A_next)
            A_next *= 0.5
            L_next = np.linalg.cholesky(A_next)
            if objective_F_cholesky(A_next,S,lam,L_next) <= init_F_val_gA + c*np.sum(g_min*A_next, dtype='float32'):
                    return A_next, step
        except linalg.LinAlgError:
            pass
        step *= beta
    return A, 0.0

def init_OBN_parser(OBN_pasrser):
    OBN_pasrser.set_defaults(algo='OBN')
    OBN_pasrser.add_argument(
        '-T', '--T', required=False, type=int, default=15, dest='T',
        help="Number of iterations.")
    OBN_pasrser.add_argument(
        '-inner_T', '--inner_T', required=False, type=int, default=50, dest='inner_T',
        help="Number of inner iterations.")
    OBN_pasrser.add_argument(
        '-linesearch', '--linesearch', required=False, type=int, default=15, dest='ls_iter',
        help="Number of linesearch iterations.")
    OBN_pasrser.add_argument(
        '-step_lim', '--step_limit', required=False, type=float, default=1e-4, dest='step_lim',
        help="The smallest step size possible.")

reason: import numpy,from numpy