text:
import sys
import math
import copy
import time
import numpy as np
import torch
from pykeops.torch import LazyTensor
from .sampling import quat_mult
from .kernels import squared_distances_tensor, lazy_interaction_kernel, \
    lazy_morse, lazy_quadratic
from .toolbox import uniform_grid_centroids, volume_ball, \
    block_sparse_reduction_parameters, maximal_eigenvector_dim2
# from pykeops.torch.cluster import cluster_ranges
from pykeops.torch.cluster import sort_clusters

# from pykeops.torch.cluster import from_matrix

"""Particles as Python iterators."""

class Particles:
    r"""The main class.

    Attributes:
        N (int): Number of particles.
        d (int): Dimension.
        pos ((N,d) Tensor): Positions of the particles.
        R (float or (N,) Tensor): Radius of the particles.
        L ((d,) Tensor): Box size.
        angle (float): Vision angle.
        axis_is_none (bool): True when an axis is specified.
        axis ((N,d) Tensor or None): Axis of the particles.
        bc (list or str): Boundary conditions. Can be
            one of the following:

            - list of size :math:`d` containing 0 (periodic) or 1
              (wall with reflecting boundary conditions) for each
              dimension.

            - ``"open"`` : no boundary conditions.

            - ``"periodic"`` : periodic boundary conditions.

            - ``"spherical"`` : reflecting boundary conditions on the
              sphere of radius ``L[0]/2`` and center **L**/2.

        target_method (dict): Dictionary of target methods (see the
            function :func:`compute_target`).
        target_option_method (dict): Dictionary of options for the
            targets (see the function :func:`compute_target`).
        blocksparse (bool): Use block sparse reduction or not.
        iteration (int): Iteration.
        centroids (Tensor): Centroids.
        eps (Tensor): Size of the cells.
        keep (BoolTensor): ``keep[i,j]`` is True when the cells ``i``
            and ``j`` are contiguous.

    """

    def __init__(self,
                 pos,
                 interaction_radius,
                 box_size,
                 vision_angle=2 * math.pi,
                 axis=None,
                 boundary_conditions='periodic',
                 block_sparse_reduction=False,
                 number_of_cells=1024):
        r"""Initialise the attributes of the class.

        Args:
            pos ((N,d) Tensor): Positions.
            interaction_radius (float or (N,) Tensor): Radius.
            box_size (float or list of size d): Box size.
            vision_angle (float, optional): Vision angle. Default is
                2*pi.
            axis ((N,d) Tensor or None, optional): Axis of the
                particles. Default is None.
            boundary_conditions (list or str, optional): Boundary conditions.
                Can be one of the following:

                - list of size :math:`d` containing 0 (periodic) or 1
                  (wall with reflecting boundary conditions) for each
                  dimension.

                - ``"open"`` : no boundary conditions.

                - ``"periodic"`` : periodic boundary conditions.

                - ``"spherical"`` : reflecting boundary conditions on the
                  sphere of radius ``L[0]/2`` and center **L**/2.

                Default is ``"periodic"``.
            block_sparse_reduction (bool, optional): Use block sparse
                reduction or not. Default is False.
            number_of_cells (int, optional): Maximum number of cells if
                **block_sparse_reduction** is True. Will be rounded to
                the nearest lower power :attr:`d`. Default is 1024.

        """
        self.N, self.d = list(pos.size())
        self.pos = pos
        self.R = interaction_radius
        self.L = torch.tensor(box_size, dtype=pos.dtype, device=pos.device)
        if self.L.shape == torch.Size([]):
            self.L = torch.repeat_interleave(self.L, self.d)
        self.angle = vision_angle
        self.axis_is_none = axis is None
        self.axis = axis
        self.bc = boundary_conditions
        self.target_method = {}
        self.add_target_method("morse", self.morse_target)
        self.add_target_method("quadratic_potential",
                               self.quadratic_potential_target)
        self.add_target_method("overlapping_repulsion",
                               self.overlapping_repulsion_target)
        self.target_option_method = {}
        self.blocksparse = block_sparse_reduction
        (self.eps,
         self.keep,
         self.centroids) = self.compute_blocksparse_parameters(number_of_cells)
        self.iteration = 0

    def __iter__(self):
        return self

    def __next__(self):
        info = self.update()
        self.iteration += 1
        return info

    def update(self):
        raise NotImplementedError()

    def compute_blocksparse_parameters(self, number_of_cells):
        r"""Compute the centroids, the size of the cells and the keep
        matrix.

        Args:
            number_of_cells (int): The maximal number of cells. It will
                be rounded to the nearest lower power :math:`d`.

        Returns:

            (nb_cells, d) Tensor, (nb_cells,) Tensor,
            (nb_cells, nb_cells) BoolTensor:

            The centroids the size of the cells and the keep matrix.

        """
        if self.blocksparse:
            if self.d == 1 or self.d > 3:
                raise NotImplementedError()
            else:
                if isinstance(self.R, float):
                    nbox = (self.L / self.R).floor()
                else:
                    nbox = (self.L / self.R.min()).floor()
                if torch.prod(nbox) <= number_of_cells:
                    eps = self.L / nbox
                    centroids = uniform_grid_centroids(eps, self.L, self.d)
                    D = squared_distances_tensor(
                        centroids,
                        centroids,
                        self.L,
                        boundary_conditions=self.bc
                    )
                    keep = D < ((1.05 * eps) ** 2).sum()
                else:
                    eps0 = (self.L.prod() / number_of_cells) ** (1. / self.d)
                    while eps0 < self.R:
                        eps0 = 1.05 * eps0
                    nbox = (self.L / eps0).floor()
                    eps = self.L / nbox
                    centroids = uniform_grid_centroids(eps, self.L, self.d)
                    D = squared_distances_tensor(
                        centroids,
                        centroids,
                        self.L,
                        boundary_conditions=self.bc
                    )
                    keep = D < ((1.05 * eps) ** 2).sum()
        else:
            eps, keep, centroids = None, None, None
        return eps, keep, centroids

    def best_blocksparse_parameters(self, min_nb_cells,
                                    max_nb_cells, step=1, nb_calls=100,
                                    loop=1, set_best=True):
        r"""Returns the optimal block sparse parameters.

        This function measures the time to execute **nb_calls**
        calls to the :meth:`__next__` for different values of the block
        sparse parameters given by the number of cells to the power
        :math:`1/d`.

        Args:
            min_nb_cells (int): Minimal number of cells to the power
                :math:`1/d`.
            max_nb_cells (int): Maximal number of cells to the power
                :math:`1/d`.
            step (int, optional): Step between two numbers of cells.
                Default is 1.
            nb_calls (int, optional): Number of calls to the function
                :meth:`__next__` at each test. Default is 100.
            loop (int, optional): Number of loops. Default is 1.
            set_best (bool, optional): Set the block sparse parameters
                to the optimal ones. Default is True.

        Returns:
            tuple:

            * **fastest** `(int)`: The number of cells which gave the
                fastest result.
            * **nb_cells** `(list)`: The number of cells tested.
            * **average_simu_time** `(list)`: The average simulation
                times.
            * **simulation_time** `(numpy array)`: The total simulation
                times.

        """
        nb_cells = list(np.arange(min_nb_cells, max_nb_cells, step))
        simulation_time = np.zeros((loop, len(nb_cells)))
        for n, nb in enumerate(nb_cells):
            for i in range(loop):
                new_copy = copy.deepcopy(self)
                new_copy.blocksparse = True
                (new_copy.eps,
                 new_copy.keep,
                 new_copy.centroids) = self.compute_blocksparse_parameters(
                    nb ** new_copy.d)
                start = time.time()
                for k in range(nb_calls):
                    new_copy.__next__()
                end = time.time()
                simulation_time[i, n] = end-start
            progress = round(100*n/len(nb_cells), 2)
            sys.stdout.write('\r' + "Progress:" + str(progress) + "%")
        average_simu_time = np.mean(simulation_time, axis=0)
        fastest = nb_cells[np.argmin(average_simu_time)]
        if set_best:
            self.blocksparse=True
            (self.eps,
             self.keep,
             self.centroids) = self.compute_blocksparse_parameters(
                fastest ** self.d)
        return fastest, nb_cells, average_simu_time, simulation_time

    def add_target_method(self, name, method):
        r"""Add a method to the dictionary :attr:`target_method`.

        Args:
            name (str): The name of the method.
            method (func): A method.

        """
        self.target_method[name] = method

    def add_target_option_method(self, name, method):
        r"""Add an option to the dictionary :attr:`target_option_method`

        Args:
            name (str): The name of the option.
            method (func): An option.

        """
        self.target_option_method[name] = method

    def compute_target(self, which_target, which_options, **kwargs):
        r"""Compute a target and apply some options.

        Args:
            which_target (dict): Dictionary of the form::

                    {"name"       : "method_name",
                     "parameters" : {"param1" : param1,
                                     "param2" : param2,
                                      ...
                                    }
                    }

                The sub-dictionary ``"parameters"`` (eventually empty)
                contains the keyword arguments to be passed to the
                function found in the dictionary :attr:`target_method`
                with the keyword "method_name".
            which_options (dict): Dictionary of the form::

                    {"name_of_option1" : parameters_of_option1,
                     "name_of_option2" : parameters_of_option2,
                    ...
                    }

                The dictionary **`parameters_of_option1`** contains
                the keyword arguments to be passed to the method found
                in the dictionary :attr:`target_option_method` with the
                keyword "name_of_option1".
            **kwargs: Keywords arguments to be passed to the functions
                which compute the target and the options.

        Returns:
            Tensor:

            Compute a target using **which_target** and
            then apply the options in **which_options**
            successively.

        """
        target = self.target_method[which_target["name"]](
            **which_target["parameters"],
            **kwargs)
        if bool(which_options):
            for option_name in which_options:
                target = self.target_option_method[option_name](
                    target,
                    **which_options[option_name],
                    **kwargs)
        return target

    def linear_local_average(self, *to_average,
                             who=None, with_who=None,
                             isaverage=True,
                             kernel=lazy_interaction_kernel):
        r"""Fast computation of a linear local average.

        A linear local average is a matrix vector product between an
        interaction kernel (LazyTensor of size (M,N)) and a Tensor
        of size (N,dimU).

        Args:
            *to_average (Tensor): The quantities to average. The first
                dimension must be equal to :attr:`N`.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            isaverage (bool, optional): Divide the result by :attr:`N`
                if **isaverage** is True. Default is True.
            kernel (func, optional): The function which returns the
                interaction kernel as a LazyTensor. Default is
                :func:`lazy_interaction_kernel`.

        Returns:
            tuple of Tensors:

            The linear local averages.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        if with_who is None:
            with_who = torch.ones(self.N,
                                  dtype=torch.bool, device=self.pos.device)

        x = self.pos[who, :]
        y = self.pos[with_who, :]
        Rx = self.R[who] if type(x) == type(self.R) else self.R
        Ry = self.R[with_who] if type(y) == type(self.R) else self.R
        if self.axis is None:
            axis = None
        else:
            axis = self.axis[who, :]

        M, D = list(x.shape)
        N, D = list(y.shape)

        U = torch.cat(to_average, dim=1)
        dimensions = []
        for u in to_average:
            dimensions.append(u.shape[1])

        dimU = U.shape[1]
        U = U[with_who, :]

        if self.blocksparse:

            (x_sorted,
             y_sorted,
             nb_centro,
             labels_x,
             labels_y,
             permutation_x,
             ranges_ij,
             map_dummies) = block_sparse_reduction_parameters(
                x, y, self.centroids, self.eps, self.L, self.keep,
                where_dummies=False)

            if type(self.R) == type(x):
                Rx_plus = torch.cat((Rx, torch.zeros(nb_centro,
                                                     dtype=Rx.dtype,
                                                     device=Rx.device)))
                Rx_perm = Rx_plus[permutation_x]
                Ry_plus = torch.cat((Ry, torch.zeros(nb_centro,
                                                     dtype=Ry.dtype,
                                                     device=Ry.device)))
                Ry_perm, lryp = sort_clusters(Ry_plus, labels_y)
            else:
                Rx_perm = Rx
                Ry_perm = Ry

            U_plus = torch.cat((U, torch.zeros(nb_centro, dimU,
                                               dtype=U.dtype,
                                               device=U.device)))
            U_sorted, labels_u = sort_clusters(U_plus, labels_y)

            if self.angle < 2 * math.pi:
                dummy_axes = torch.randn((nb_centro, self.d),
                                         dtype = axis.dtype,
                                         device=axis.device)
                dummy_axes = dummy_axes / torch.norm(
                    dummy_axes, dim=1).reshape((nb_centro, 1))
                axis = torch.cat((axis, dummy_axes))
                axis = axis[permutation_x, :]

            K_ij = kernel(x=x_sorted, y=y_sorted,
                          Rx=Rx_perm, Ry=Ry_perm,
                          L=self.L, boundary_conditions=self.bc,
                          vision_angle=self.angle, axis=axis)

            K_ij.ranges = ranges_ij

            J_perm = K_ij @ U_sorted
            J_plus = J_perm[torch.argsort(permutation_x), :]
            # Remove the dummy particles
            J = J_plus[0:M, :]
            if isaverage:
                J = (1. / N) * J

        else:
            K_ij = kernel(x=x, y=y,
                          Rx=Rx, Ry=Ry,
                          L=self.L, boundary_conditions=self.bc,
                          vision_angle=self.angle, axis=axis)

            J = K_ij @ U
            if isaverage:
                J = (1. / self.N) * J
        return torch.split(J, dimensions, dim=1)

    def nonlinear_local_average(self, binary_formula,
                                arg1, arg2,
                                who=None, with_who=None,
                                isaverage=True,
                                kernel=lazy_interaction_kernel):
        r"""Fast computation of a nonlinear local average.

        A nonlinear local average is a reduction of the form

        .. math::
            \sum_j K_{ij} U_{ij}

        where :math:`K_{ij}` is an interaction kernel (
        LazyTensor of size (M,N)) and :math:`U_{ij}` is a
        LazyTensor of size (M,N,dim) computed using a binary
        formula.

        Args:
            binary_formula (func): Takes two arguments **arg1**
                and **arg2** and returns a LazyTensor of size
                (M,N,dim).
            arg1 ((M,D1) Tensor): First argument of the binary formula.
            arg2 ((N,D2) Tensor): Second argument of the binary
                formula.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            isaverage (bool, optional): Divide the result by :attr:`N`
                if isaverage is True. Default is True.
            kernel (func, optional): The function which returns the
                interaction kernel as a LazyTensor. Default is
                :func:`lazy_interaction_kernel`.

        Returns:
            (M,dim) Tensor:

            The nonlinear local average.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        if with_who is None:
            with_who = torch.ones(self.N,
                                  dtype=torch.bool, device=self.pos.device)

        x = self.pos[who, :]
        y = self.pos[with_who, :]
        Rx = self.R[who] if type(x) == type(self.R) else self.R
        Ry = self.R[with_who] if type(y) == type(self.R) else self.R
        if self.axis is None:
            axis = None
        else:
            axis = self.axis[who, :]

        M, D = list(x.shape)
        N, D = list(y.shape)

        if self.blocksparse:

            (x_sorted,
             y_sorted,
             nb_centro,
             labels_x,
             labels_y,
             permutation_x,
             ranges_ij,
             map_dummies) = block_sparse_reduction_parameters(
                x, y, self.centroids, self.eps, self.L, self.keep,
                where_dummies=True)

            if type(self.R) == type(x):
                Rx_plus = torch.cat((Rx, torch.zeros(nb_centro,
                                                     dtype=Rx.dtype,
                                                     device=Rx.device)))
                Rx_perm = Rx_plus[permutation_x]
                Ry_plus = torch.cat((Ry, torch.zeros(nb_centro,
                                                     dtype=Ry.dtype,
                                                     device=Ry.device)))
                Ry_perm, lryp = sort_clusters(Ry_plus, labels_y)
            else:
                Rx_perm = Rx
                Ry_perm = Ry

            if self.angle < 2 * math.pi:
                dummy_axes = torch.randn((nb_centro, self.d),
                                         dtype = axis.dtype,
                                         device=axis.device)
                dummy_axes = dummy_axes / torch.norm(
                    dummy_axes, dim=1).reshape((nb_centro, 1))
                axis = torch.cat((axis, dummy_axes))
                axis = axis[permutation_x, :]

            K_ij = kernel(x=x_sorted, y=y_sorted,
                          Rx=Rx_perm, Ry=Ry_perm,
                          L=self.L, boundary_conditions=self.bc,
                          vision_angle=self.angle, axis=axis)

            N1, D1 = list(arg1.size())
            N2, D2 = list(arg2.size())
            arg1_plus = torch.cat((arg1, torch.zeros(nb_centro, D1,
                                                     dtype=arg1.dtype,
                                                     device=arg1.device)))
            arg2_plus = torch.cat((arg2, torch.zeros(nb_centro, D2,
                                                     dtype=arg2.dtype,
                                                     device=arg2.device)))
            arg1_plus_sorted, la1 = sort_clusters(arg1_plus, labels_x)
            arg2_plus_sorted, la2 = sort_clusters(arg2_plus, labels_y)
            U_ij_plus = map_dummies * binary_formula(arg1_plus_sorted,
                                                     arg2_plus_sorted)

            KU = K_ij * U_ij_plus
            KU.ranges = ranges_ij
            J_perm = KU.sum(1)
            J_plus = J_perm[torch.argsort(permutation_x), :]
            # Remove the dummy particles
            J = J_plus[0:M, :]
            if isaverage:
                J = (1. / N) * J

        else:
            K_ij = kernel(x=x, y=y,
                          Rx=Rx, Ry=Ry,
                          L=self.L, boundary_conditions=self.bc,
                          vision_angle=self.angle, axis=axis)
            U_ij = binary_formula(arg1, arg2)
            KU = K_ij * U_ij
            J = KU.sum(1)
            if isaverage:
                J = (1. / N) * J

        return J

    def number_of_neighbours(self, who=None, with_who=None):
        r""" Compute the number of neighbours.

        The number of neighbours of ``x[i,:]`` is the number of ones in
        the row corresponding to ``x[i,:]`` in the interaction kernel.

        Args:
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.

        Returns:
            (M,) Tensor:

            The number of neighbours.

        """
        Nneigh, = self.linear_local_average(
            torch.ones((self.N, 1),
                       dtype=self.pos.dtype, device=self.pos.device),
            who=who, with_who=with_who, isaverage=False)
        return Nneigh.reshape(Nneigh.shape[0])

    def morse_target(self, Ca, la, Cr, lr, p=2, mass=1.,
                     who=None, with_who=None,
                     local=False,
                     isaverage=True,
                     kernel=lazy_interaction_kernel, **kwargs):
        r"""Compute the force exerted on each particle due to the Morse
        potential.

        Args:
            Ca (float): Attraction coefficient.
            la (float): Attraction length.
            Cr (float): Repulsion coefficient.
            lr (float): Repulsion length.
            p (int, optional): Exponent. Default is 2.
            mass ((N,) Tensor or float, optional): Mass of the
                particles. Default is 1.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            local (bool, optional): If **local** is true then the sum
                only considers the y particles within the interaction
                kernel. Default is False.
            isaverage (bool, optional): If isaverage is True then divide
                the total force by N where N is the number of True in
                :Tensor"`with_who`. Default is True.
            kernel (func, optional): The interaction kernel to use.
                Default is :func:`lazy_interaction_kernel`
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The force exerted on each particle located at ``pos[who,:]``.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        if with_who is None:
            with_who = torch.ones(self.N,
                                  dtype=torch.bool, device=self.pos.device)
        x = self.pos[who, :]
        y = self.pos[with_who, :]
        mx = mass[who] if type(x) == type(mass) else mass
        my = mass[with_who] if type(y) == type(mass) else mass
        if not local:
            lazy_morse_tensor = lazy_morse(x=x, y=y, Ca=Ca, la=la,
                                           Cr=Cr, lr=lr, p=p, mx=mx, my=my)
            if isaverage:
                return (1. / y.shape[0]) * lazy_morse_tensor.sum(axis=1)
            else:
                return lazy_morse_tensor.sum(axis=1)
        else:
            binary_formula = lambda x, y: lazy_morse(x=x, y=y, Ca=Ca, la=la,
                                                     Cr=Cr, lr=lr, p=p,
                                                     mx=mx, my=my)
            return self.nonlinear_local_average(
                binary_formula, x, y,
                who=who, with_who=with_who,
                isaverage=isaverage,
                kernel=kernel)

    def quadratic_potential_target(self, who=None, with_who=None,
                                   kernel=lazy_interaction_kernel,
                                   isaverage=True, **kwargs):
        r"""Compute the force exerted on each particle due to the
        quadratic potential.

        Args:
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            isaverage (bool, optional): If isaverage is True then divide
                the total force by N where N is the number of True in
                :Tensor"`with_who`. Default is True.
            kernel (func, optional): The interaction kernel to use.
                Default is :func:`lazy_interaction_kernel`
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The force exerted on each particle located at ``pos[who,:]``.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        if with_who is None:
            with_who = torch.ones(self.N,
                                  dtype=torch.bool, device=self.pos.device)
        x = self.pos[who, :]
        y = self.pos[with_who, :]
        binary_formula = lambda x, y: lazy_quadratic(
            x=x, y=y, R=self.R, L=self.L, boundary_conditions=self.bc)
        return self.nonlinear_local_average(
            binary_formula, x, y,
            who=who, with_who=with_who,
            isaverage=isaverage,
            kernel=kernel)

    def overlapping_repulsion_target(self, who=None, with_who=None, **kwargs):
        r"""Compute the repulsion force exerted on each particle due to
        the overlapping.

        The force exerted on a particle located at :math:`x_i` derives
        from a logarithmic potential and is given by:

        .. math::

            F = \sum_j K(x_i-x_j) \left(\frac{(R_i+R_j)^2}{|x_i-x_j|^2} - 1\right) (x_i-x_j)

        where :math:`x_j` is the position of the other particles;
        :math:`R_i` and :math:`R_j` are the radii of the particles at
        :math:`x_i` and :math:`x_j` respectively and :math:`K` is the
        overlapping kernel.

        Note:
            in the present case, it is simpler to write explicitly
            the reduction formula but the same function could be
            implemented using a binary formula in
            :func:`nonlinear_local_average` with the kernel given by
            :func:`lazy_overlapping_kernel`.

        Args:
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The force exerted on each particle located at ``pos[who,:]``.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        if with_who is None:
            with_who = torch.ones(self.N,
                                  dtype=torch.bool, device=self.pos.device)
        x = self.pos[who, :]  # (M,D)
        y = self.pos[with_who, :]  # (N,D)
        Rx = self.R[who] if type(self.R) == type(x) else self.R
        Ry = self.R[with_who] if type(self.R) == type(y) else self.R

        x_i = LazyTensor(x[:, None, :])
        y_j = LazyTensor(y[None, :, :])
        if type(Rx) == type(x):
            R_i = LazyTensor(Rx.reshape((x.shape[0], 1))[:, None])
            R_j = LazyTensor(Ry.reshape((y.shape[0], 1))[None, :])
            R_ij = R_i + R_j  # (M,N)
        else:
            R_ij = Rx + Ry  # float
        z_ij = x_i - y_j  # (M,N,D)
        sq_dist = (z_ij ** 2).sum(-1)  # (M,N)
        R_ij = R_ij * sq_dist.sign()
        sq_dist = sq_dist - (sq_dist.sign() - 1)
        F_ij = (((R_ij ** 2) / sq_dist - 1).relu()) * z_ij

        return F_ij.sum(axis=1)

"""Subclass of kinetic particles."""

class KineticParticles(Particles):
    r"""Subclass of kinetic particles.

    Attributes:
         vel ((N,d) Tensor): Velocities.

    """

    def __init__(self,
                 pos,
                 vel,
                 interaction_radius,
                 box_size,
                 vision_angle=2 * math.pi,
                 axis=None,
                 boundary_conditions='periodic',
                 block_sparse_reduction=False,
                 number_of_cells=1024):
        r"""

        Args:
            pos ((N,d) Tensor): Positions.
            vel ((N,d) Tensor): Velocities.
            interaction_radius (float or (N,) Tensor): Radius.
            box_size (float or list of size d): Box size.
            vision_angle (float, optional): Vision angle. Default is
                :math:`2\pi`.
            axis ((N,d) Tensor or None, optional): Axis of the
                particles. Default is None.
            boundary_conditions (list or str): Boundary conditions.
                Can be one of the following:

                - list of size :math:`d` containing 0 (periodic) or 1
                  (wall with reflecting boundary conditions) for each
                  dimension.

                - ``"open"`` : no boundary conditions.

                - ``"periodic"`` : periodic boundary conditions.

                - ``"spherical"`` : reflecting boundary conditions on the
                  sphere of radius ``L[0]/2`` and center **L**/2.
                Default is ``"periodic"``.
            block_sparse_reduction (bool, optional): Use block sparse
                reduction or not. Default is False.
            number_of_cells (int, optional): Maximum number of cells if
                **block_sparse_reduction** is True. Will be rounded to
                the nearest lower power :attr:`d`. Default is 1024.

        """
        self.vel = vel
        super().__init__(pos=pos,
                         interaction_radius=interaction_radius,
                         box_size=box_size,
                         vision_angle=vision_angle,
                         axis=axis,
                         boundary_conditions=boundary_conditions,
                         block_sparse_reduction=block_sparse_reduction,
                         number_of_cells=number_of_cells)
        self.add_target_method("normalised", self.normalised)
        self.add_target_method("motsch_tadmor", self.motsch_tadmor)
        self.add_target_method("mean_field", self.mean_field)
        self.add_target_method("max_kappa", self.max_kappa)
        self.add_target_method("nematic", self.nematic)
        self.add_target_option_method("bounded_angular_velocity",
                                      self.bounded_angular_velocity)
        self.add_target_option_method("pseudo_nematic", self.pseudo_nematic)

    @property
    def axis(self):
        r"""
        If a None axis is given, then the default axis is the
        normalised velocity.

        Returns:
            (N,d) Tensor

        """
        if self.axis_is_none:
            return self.vel / torch.norm(self.vel, dim=1).reshape((self.N, 1))
        else:
            return self.__axis

    @axis.setter
    def axis(self, value):
        if self.axis_is_none:
            """DO NOTHING"""
        else:
            self.__axis = value

    @property
    def order_parameter(self):
        r"""Compute the norm of the average velocity.

        Returns:
            float:

            The order parameter.

        """
        vec_phi = (1. / self.N) * self.vel.sum(0)
        phi = torch.norm(vec_phi)
        return float(phi)

    def check_boundary(self):
        r"""Update the positions and velocities of the particles which
           are outside the boundary.

        """
        if self.bc == 'periodic':
            for i in range(self.d):
                self.pos[:, i] = torch.remainder(self.pos[:, i], self.L[i])
        elif type(self.bc) == list:
            for i in range(self.d):
                if self.bc[i] == 1:
                    inf0 = self.pos[:, i] < 0
                    supL = self.pos[:, i] > self.L[i]
                    self.pos[inf0, i] = -self.pos[inf0, i]
                    self.pos[supL, i] = 2 * self.L[i] - self.pos[supL, i]

                    self.vel[inf0, i] = -self.vel[inf0, i]
                    self.vel[supL, i] = -self.vel[supL, i]
                elif self.bc[i] == 0:
                    self.pos[:, i] = torch.remainder(self.pos[:, i], self.L[i])
        elif self.bc == 'spherical':
            # Translate the center to the origin and compute the
            # radius = L/2.
            x = self.pos - self.L / 2
            r = self.L[0] / 2
            out = (x ** 2).sum(1) > (r ** 2)
            N_out = out.sum()
            if N_out > 0:
                x_out = x[out, :]
                v_out = self.vel[out, :]
                # Find x0 and c0 such that x = x0 + c0*vel with |x0|=r
                xdotv = (x_out * v_out).sum(1)
                sqnorm_v_out = (v_out ** 2).sum(1)
                sqnorm_x_out = (x_out ** 2).sum(1)
                D0 = (xdotv ** 2
                      - sqnorm_v_out * (sqnorm_x_out - r ** 2)
                      ).abs()  # Take the absolute value to ensure >=0
                c0 = 1. / sqnorm_v_out * (xdotv - D0.sqrt())
                c0 = c0.reshape((N_out, 1))
                x0 = x_out - c0 * v_out
                # Find x1 and c1 such that x = x1+c1*x0 with |x1|=r
                xdotx0 = (x_out * x0).sum(1)
                D1 = (xdotx0 ** 2 - (r ** 2) * (sqnorm_x_out - r ** 2)).abs()
                c1 = (r ** (-2)) * (xdotx0 - D1.sqrt())
                c1 = c1.reshape((N_out, 1))
                # Find c2 such that v = v2+c2*x0 and |v|=|v2|
                c2 = 2. * (r ** (-2)) * ((v_out * x0).sum(1))
                c2 = c2.reshape((N_out, 1))
                # Construct the reflections
                x_reflect = x_out - 2 * c1 * x0
                v_reflect = v_out - c2 * x0
                v_reflect = v_reflect / torch.norm(v_reflect,
                                                   dim=1).reshape((N_out, 1)) \
                            * sqnorm_v_out.sqrt().reshape((N_out, 1))
                # Update the positions and velocity
                self.pos[out, :] = x_reflect + self.L / 2
                self.vel[out, :] = v_reflect
                # Stability and emergency procedure....
                self.pos = torch.max(self.pos,
                                     torch.tensor([0.],
                                                  dtype=self.pos.dtype,
                                                  device=self.pos.device))
                self.pos = torch.min(self.pos, self.L[0])
                x = self.pos - self.L / 2
                r = self.L[0] / 2
                out = (x ** 2).sum(1) > (r ** 2)
                N_out = out.sum()
                if N_out > 0:
                    if N_out > 42:
                        print('Warning!'+str(N_out)+'particles out')
                    self.pos[out, :] = self.L/2 + torch.zeros(
                        (N_out, self.d),
                        dtype=self.pos.dtype,
                        device=self.pos.device
                    )

    def normalised(self, kappa, who=None, with_who=None, **kwargs):
        r"""The normalised average velocity.

        Given a system of particles with positions :math:`(x_i)_i` and
        velocities :math:`(v_i)_i`, the normalised average velocity
        around particle :math:`(x_i,v_i)` is

        .. math::

            \Omega_i = \frac{\sum_j K(x_i-x_j) v_j}{\left|\sum_j K(x_i-x_j) v_j\right|},

        where :math:`K` is the interaction kernel.

        Args:
            kappa ((1,) Tensor or (M,) Tensor): The concentration
                parameter. If the size of kappa is bigger than 1 then
                its size must be equal to the number of True values in
                **who**.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The normalised averaged velocity multiplied by the
            concentration parameter.

        """
        J, = self.linear_local_average(self.vel, who=who, with_who=with_who)
        Omega = J / torch.norm(J, dim=1).reshape((J.shape[0], 1))
        return kappa * Omega

    def motsch_tadmor(self, kappa, who=None, with_who=None, **kwargs):
        r"""The average velocity of the neighbours only.

        Given a system of particles with positions :math:`(x_i)_i` and
        velocities :math:`(v_i)_i`, the average velocity
        around particle :math:`(x_i,v_i)` with the Motsch-Tadmor
        normalisation is:

        .. math::

            \Omega_i = \frac{\sum_j K(x_i-x_j) v_j}{\sum_j K(x_i-x_j)},

        where :math:`K` is the interaction kernel.

        Args:
            kappa ((1,) Tensor or (M,) Tensor): The concentration
                parameter. If the size of kappa is bigger than 1 then
                its size must be equal to the number of True values in
                **who**.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The average velocity of the neighbours multiplied
            by the concentration parameter.

        """
        justN = torch.ones((self.N, 1),
                           dtype=self.pos.dtype, device=self.pos.device)
        J, N_neighbours = self.linear_local_average(self.vel, justN,
                                                    who=who, with_who=with_who)
        return kappa * J / N_neighbours

    def mean_field(self, kappa, who=None, with_who=None, **kwargs):
        r"""Non-normalised average velocity.

        Given a system of particles with positions :math:`(x_i)_i` and
        velocities :math:`(v_i)_i`, the mean-field average velocity
        around particle :math:`(x_i,v_i)` (without normalisation) is:

        .. math::

            J_i = \frac{1}{N}\sum_j K(x_i-x_j) v_j,

        where :math:`K` is the interaction kernel.

        Args:
            kappa ((1,) Tensor or (M,) Tensor): The concentration
                parameter. If the size of kappa is bigger than 1 then
                its size must be equal to the number of True values in
                **who**.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The mean-field average velocity multiplied by a scaling
            parameter and by the concentration parameter. The scaling
            parameter is equal to the volume of the box divided by
            the volume of the ball of radius :attr:`R` in dimension
            :attr:`d`.

        """
        scaling = (self.L / self.R) ** self.d * 1. / volume_ball(self.d)
        J, = self.linear_local_average(self.vel, who=who, with_who=with_who)
        J = scaling * J
        return kappa * J

    def max_kappa(self, kappa, kappa_max, who=None, with_who=None, **kwargs):
        r""" The mean-field target with a norm threshold.

        Given a system of particles with positions :math:`(x_i)_i` and
        velocities :math:`(v_i)_i`, the mean-field average velocity
        around particle :math:`(x_i,v_i)` with a cutoff is:

        .. math::

            \Omega_i = \frac{1}{\frac{1}{\kappa}+\frac{|J_i|}{\kappa_\mathrm{max}}} J_i

        where

        .. math::

            J_i = \frac{1}{N}\sum_j K(x_i-x_j) v_j,

        and :math:`K` is the interaction kernel multiplied by the
        scaling parameter. The scaling parameter is equal to the
        volume of the box divided by the volume of the ball of radius
        :attr:`R` in dimension :attr:`d`. The norm of :math:`\Omega_i`
        is thus bounded by :math:`\kappa_\mathrm{max}`.

        Args:
            kappa ((1,) Tensor or (M,) Tensor): The concentration
                parameter. If the size of kappa is bigger than 1 then
                its size must be equal to the number of True values in
                **who**.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The cutoff average velocity.

        """
        scaling = (self.L / self.R) ** self.d * 1. / volume_ball(self.d)
        J, = self.linear_local_average(self.vel, who=who, with_who=with_who)
        J = scaling * J
        normJ = torch.norm(J, dim=1).reshape((J.shape[0], 1))
        return 1. / ((1. / kappa) + (normJ / kappa_max)) * J

    def nematic(self, kappa, who=None, with_who=None, **kwargs):
        r"""Normalised average of the nematic velocities.

        Given a system of particles with positions :math:`(x_i)_i` and
        velocities :math:`(v_i)_i`, the nematic average velocity
        around particle :math:`(x_i,v_i)` is:

        .. math::

            \Omega_i = (v_i\cdot \overline{\Omega}_i)\overline{\Omega}_i,

        where :math:`\overline{\Omega}_i` is any unit eigenvector
        associated to the maximal eigenvalue of the average Q-tensor:

        .. math::

            Q_i = \sum_j K(x_i-x_j) \left(v_j\otimes v_j - \frac{1}{d} I_d\right),

        where :math:`K` is the interaction kernel.

        Args:
            kappa ((1,) Tensor or (M,) Tensor): The concentration
                parameter. If the size of kappa is bigger than 1 then
                its size must be equal to the number of True values in
                **who**.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The nematic average velocity multiplied by the concentration
            parameter.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        Q = torch.einsum('ijl,ilk->ijk',
                         self.vel[:, :, None], self.vel[:, None, :]) \
            - (1. / self.d) * torch.eye(self.d,
                                        dtype=self.vel.dtype,
                                        device=self.vel.device)
        reshapeQ = Q.reshape(self.N, self.d ** 2)
        J, = self.linear_local_average(reshapeQ, who=who, with_who=with_who)
        M = J.shape[0]
        J = J.reshape(M, self.d, self.d)
        if self.d == 2:
            targets = maximal_eigenvector_dim2(J)
        else:
            eivel, eivec = torch.linalg.eigh(J.to('cpu'))
            targets = eivec[:, :, -1]
            targets = targets.to(J.device)
        dot = (self.vel[who, :] * targets).sum(1).reshape((M, 1))
        targets = dot * targets
        return kappa * targets

    def bounded_angular_velocity(self, target, angvel_max, dt,
                                 who=None, **kwargs):
        r"""Bounded angle between the velocity and the target.

        The maximum angle between the velocity and the new target cannot
        exceed the value ``angvel_max * dt``. The output tensor
        **new_target** is of the form::

            new_target = a * vel + b * target

        where :math:`a,b>0` are such that the cosine between
        **new_target** and :attr:`vel` is equal to::

            max_cos := max(cos(angvel_max * dt), cos_target)

        where **cos_target** is the cosine between **target** and
        :attr:`vel`. The solution is::

            b = (1 - max_cos ** 2) / (1 - cos_target ** 2)
            a = |target| / |vel| * (max_cos - b * cos_target)

        Args:
            target ((M,d) Tensor): Target.
            angvel_max (float): Maximum angular velocity.
            dt (float): Interval of time.
            who ((N,) BoolTensor, optional): The velocities of the
                particles associated to the target are ``vel[who,:]``
                (or :attr:`vel` if **who** is None). Default is None.
            **kwargs: Arbitrary keywords arguments.

        Returns:
            (M,d) Tensor:

            The new target.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        M = target.shape[0]
        vel = self.vel[who, :]
        norm_vel = torch.norm(vel, dim=1).reshape((M, 1))
        norm_target = torch.norm(target, dim=1).reshape((M, 1))

        cos_max = math.cos(angvel_max * dt)
        dot = (vel * target).sum(axis=1).reshape((M, 1))
        cos = dot / (norm_vel * norm_target)
        max_cos = dot / (norm_vel * norm_target)
        max_cos[max_cos <= cos_max] = cos_max

        # Add .01 for stability
        b = torch.sqrt((1.01 - max_cos ** 2) / (1.01 - cos ** 2))
        a = (norm_target / norm_vel) * (max_cos - b * cos)

        new_target = a * vel + b * target

        return new_target

    def pseudo_nematic(self, target, who=None, **kwargs):
        r"""Choose the closest to the velocity between the
        target and its opposite.

        Args:
            target ((M,d) Tensor): Target.
            who ((N,) BoolTensor, optional): The velocities of the
                particles associated to the target are ``vel[who,:]``
                (or :attr:`vel` if **who** is None). Default is None.
            **kwargs: Arbitrary keywords arguments.

        Returns:
            (M,d) Tensor:

            Take the opposite of **target** if the angle with the
            velocity is larger than :math:`\pi/2`.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        dot = (self.vel[who, :] * target).sum(axis=1)
        target[dot < 0] = -target[dot < 0]
        return target

"""The subclass of body-oriented particles."""

class BOParticles(Particles):
    r"""3D particle with a body-orientation in :math:`SO(3)`.

    A body-orientation is a rotation matrix is :math:`SO(3)` or
    equivalently a unit quaternion.
    If :math:`q = (x, y, z, t)` is a unit quaternion, the associated
    rotation matrix is :math:`A = [e_1, e_2, e_3]` where the column
    vectors are defined by:

    .. math::

        e_1 = (x^2 + y^2 - z^2 + t^2, 2(xt+yz), 2(yt-xz))^T

    .. math::
        e_2 = (2(yz-xt), x^2 - y^2 + z^2 - t^2, 2(xy+zt))^T

    .. math::
        e_3 = (2(xz+yt), 2(zt-xy), x^2 - y^2 - z^2 + t^2)^T

    Conversely, if :math:`A` is the rotation of angle :math:`\theta`
    around the axis :math:`u = (u_1, u_2, u_3)^T` then the associated
    unit quaternion is:

    .. math::

        q = \cos(\theta/2) + \sin(\theta/2)(iu_1 + ju_2 + ku_3).

    or its opposite.

    Attributes:
        bo ((N,4) Tensor): Body-orientation stored as a unit quaternion.

    """

    def __init__(self, pos, bo,
                 interaction_radius, box_size,
                 vision_angle=2 * math.pi, axis=None,
                 boundary_conditions='periodic',
                 block_sparse_reduction=False, number_of_cells=1024):
        r"""

        Args:
            pos ((N,3) Tensor): Positions.
            bo ((N,4) Tensor): Body-orientations (unit quaternions).
            interaction_radius (float or (N,) Tensor): Radius.
            box_size (float or list of size d): Box size.
            vision_angle (float, optional): Vision angle. Default is
                :math:`2\pi`.
            axis ((N,3) Tensor or None, optional): Axis of the
                particles. Default is None.
            boundary_conditions (list or str, optional): Boundary conditions.
                Can be one of the following:

                - list of size :math:`d` containing 0 (periodic) or 1
                  (wall with reflecting boundary conditions) for each
                  dimension.

                - ``"open"`` : no boundary conditions.

                - ``"periodic"`` : periodic boundary conditions.

                - ``"spherical"`` : reflecting boundary conditions on the
                  sphere of radius ``L[0]/2`` and center **L**/2.

                Default is ``"periodic"``.
            block_sparse_reduction (bool, optional): Use block sparse
                reduction or not. Default is False
            number_of_cells (int, optional): Maximum number of cells if
                **block_sparse_reduction** is True. Will be rounded to
                the nearest lower power :attr:`d`. Default is 1024.

        """
        self.bo = bo
        super().__init__(pos,
                         interaction_radius,
                         box_size,
                         vision_angle=vision_angle, axis=axis,
                         boundary_conditions=boundary_conditions,
                         block_sparse_reduction=block_sparse_reduction,
                         number_of_cells=number_of_cells)
        self.add_target_method("normalised", self.normalised)

    @property
    def vel(self):
        r"""The velocity is the first column vector of the
        body-orientation matrix.

        Returns:
            (N,3) Tensor:

            The velocity.

        """
        x = self.bo[:, 0]
        y = self.bo[:, 1]
        z = self.bo[:, 2]
        t = self.bo[:, 3]
        e1 = torch.stack(
            (x ** 2 + y ** 2 - z ** 2 - t ** 2,
             2 * (x * t + y * z),
             2 * (y * t - x * z)), dim=1)
        return e1

    @property
    def axis(self):
        r""" If a None axis is given, then the default axis is the
        (normalised) velocity.

        Returns:
            (N,3) Tensor:

        """
        if self.axis_is_none:
            return self.vel
        else:
            return self.__axis

    @axis.setter
    def axis(self, value):
        if self.axis_is_none:
            """DO NOTHING"""
        else:
            self.__axis = value

    @property
    def orth_basis(self):
        r"""An othogonal basis of the orthogonal plane of the velocity.

        Returns:
            tuple:

            * **e2** *((N,3) Tensor)*: Second column vector of the
              body-orientation matrix.

            * **e3** *((N,3) Tensor)*: Third column vector of the
              body-orientation matrix.

        """
        x = self.bo[:, 0]
        y = self.bo[:, 1]
        z = self.bo[:, 2]
        t = self.bo[:, 3]
        e2 = torch.stack(
            (2 * (y * z - x * t),
             x ** 2 - y ** 2 + z ** 2 - t ** 2,
             2 * (x * y + z * t)), dim=1)
        e3 = torch.stack(
            (2 * (x * z + y * t),
             2 * (z * t - x * y),
             x ** 2 - y ** 2 - z ** 2 + t ** 2), dim=1)
        return e2, e3

    @property
    def qtensor(self):
        r"""Compute the normalised uniaxial Q-tensors of the particles.

        If :math:`q` is a unit quaternion seen as a colum vector of
        dimension 4, the normalised uniaxial Q-tensor is defined by:

        .. math::

            Q = q \otimes q - \frac{1}{4}I_4

        where :math:`\otimes` is the tensor product and :math:`I_4` the
        identity matrix of size 4.

        Returns:
            (N,4,4) Tensor:

            Q-tensors of the N particles.

        """
        Q = torch.einsum('ijl,ilk->ijk',
                         self.bo[:, :, None], self.bo[:, None, :]) \
            - .25 * torch.eye(4, dtype=self.bo.dtype, device=self.bo.device)
        return Q

    @property
    def omega_basis(self):
        r"""Local frame associated to the direction of motion.

        The unit vector :math:`\Omega` is the direction of motion
        (it is equal to :attr:`vel`) and
        :math:`[\Omega, e_\theta, e_\varphi]` is the local orthomormal
        frame associated to the spherical cooordinates.

        Returns:
            tuple:

            * **p** *((N,3) Tensor)*: :math:`e_\varphi`.
            * **q** *((N,3) Tensor)*: :math:`-e_\theta`.
        """
        Omega = self.vel
        phi = torch.atan2(Omega[:, 1], Omega[:, 0])
        p = torch.cat(
            (-torch.sin(phi).reshape((self.N, 1)),
             torch.cos(phi).reshape((self.N, 1)),
             torch.zeros((self.N, 1), dtype=phi.dtype, device=phi.device)),
            dim=1)
        q = torch.cross(Omega, p)
        return p, q

    @property
    def u_in_omega_basis(self):
        r"""Coordinates of :math:`u` in the frame :math:`[\Omega, p, q]`
        given by :attr:`omega_basis` where :math:`u` is the second
        column vector of the body-orientation matrix."""
        p, q = self.omega_basis
        u, v = self.orth_basis
        up = (u * p).sum(dim=1)
        uq = (u * q).sum(dim=1)
        return up, uq

    def normalised(self, who=None, with_who=None, **kwargs):
        r"""Projection of the average body-orientation on :math:`SO(3)`.

        Given a system of particles with positions :math:`(x_i)_i` and
        body-orientation matrices :math:`(A_i)_i`, the mean-field average
        body-orientation around the particle :math:`(x_i,A_i)` is
        is the arithmetic average:

        .. math::

            J_i = \sum_j K(x_i - x_j) A_j

        where :math:`K` is the interaction kernel. The projection on
        :math:`SO(3)` is

        .. math::

            P_{SO(3)}(J_i) = argmax_{A\in SO(3)}( A \cdot J_i )

        for the dot product

        .. math::

            A \cdot B = \frac{1}{2} Tr(A^T B).

        The unit quaternion associated to :math:`P_{SO(3)}(J_i)` is any
        eigenvector associated to the maximal eigenvalue of the average
        Q-tensor:

        .. math::

            \overline{Q}_i = \sum_j K(x_i - x_j) Q_j

        where :math:`Q_j` is the Q-tensor of particle :math:`(x_j,A_j)`
        given by :attr:`qtensor`.

        Args:
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            (M,4) Tensor:

            The unit quaternion associated to the projection of the
            mean-field average body-orientation on :math:`SO(3)`.

        """
        reshapeQ = self.qtensor.reshape(self.N, 16)
        J, = self.linear_local_average(reshapeQ, who=who, with_who=with_who)
        M = J.shape[0]
        J = J.reshape(M, 4, 4)
        # Faster on the cpu !
        eivel, eivec = torch.linalg.eigh(J.to('cpu'))
        targets = eivec[:, :, 3]
        # Back to the gpu.
        targets = targets.to(J.device)
        return targets

    ### Some statistical quantities ###

    @property
    def theta(self):
        r"""Polar angle of the direction of motion"""
        Omega = self.vel.sum(0)
        Omega = Omega / torch.norm(Omega)
        return float(torch.acos(Omega[2]))

    @property
    def phi(self):
        r"""Azimuthal angle of the direction of motion"""
        Omega = self.vel.sum(0)
        Omega = Omega / torch.norm(Omega)
        phi_pi = float(torch.atan2(Omega[1], Omega[0]))
        return np.mod(phi_pi, 2 * np.pi)

    def local_vel(self, xmin=0, xmax=1, ymin=0, ymax=1, zmin=0, zmax=1,
                  who=None):
        if who is None:
            index = (self.pos[:, 0] >= xmin and self.pos[:, 0] <= xmax) \
                    and (self.pos[:, 0] >= xmin and self.pos[:, 0] <= xmax) \
                    and (self.pos[:, 0] >= xmin and self.pos[:, 0] <= xmax)
        else:
            index = who
        Omega = self.vel[index, :].sum(0)
        order_parameter = torch.norm(Omega) / len(self.vel[index, 0])
        Omega = Omega / torch.norm(Omega)
        theta = float(torch.acos(Omega[2]))
        phi = np.mod(float(torch.atan2(Omega[1], Omega[0])), 2 * np.pi)
        return order_parameter, theta, phi

    def twist(self, K, axis, return_slice_nb=None, local_vel=False):
        up, uq = self.u_in_omega_basis
        if axis == 'x':
            if self.L.shape == torch.Size([]):
                eps = self.L / K
            else:
                eps = self.L[0] / K
            index = (self.pos[:, 0] / eps).floor().int()
        elif axis == 'y':
            if self.L.shape == torch.Size([]):
                eps = self.L / K
            else:
                eps = self.L[1] / K
            index = (self.pos[:, 1] / eps).floor().int()
        elif axis == 'z':
            if self.L.shape == torch.Size([]):
                eps = self.L / K
            else:
                eps = self.L[2] / K
            index = (self.pos[:, 2] / eps).floor().int()
        else:
            raise ValueError('axis must be x, y or z')
        index = torch.max(index, torch.tensor(0, device=index.device))
        index = torch.min(index, torch.tensor(K - 1, device=index.device))
        slice_p = torch.zeros(K, dtype=up.dtype, device=up.device)
        slice_q = torch.zeros(K, dtype=uq.dtype, device=uq.device)
        if local_vel:
            op = torch.zeros(K, dtype=up.dtype, device=up.device)
            theta_loc = torch.zeros(K, dtype=up.dtype, device=up.device)
            phi_loc = torch.zeros(K, dtype=up.dtype, device=up.device)
        for k in range(K):
            Nk = torch.sum(index == k)
            slice_p[k] = torch.sum(up[index == k]) / Nk
            slice_q[k] = torch.sum(uq[index == k]) / Nk
            if local_vel:
                who_k = (index == k)
                opk, thk, phk = self.local_vel(who=who_k)
                op[k] = opk
                theta_loc[k] = thk
                phi_loc[k] = phk
        if return_slice_nb is None:
            if local_vel:
                return slice_p, slice_q, op, theta_loc, phi_loc
            else:
                return slice_p, slice_q
        else:
            if local_vel:
                return (slice_p, slice_q,
                        up[index == return_slice_nb],
                        uq[index == return_slice_nb],
                        op, theta_loc, phi_loc)
            else:
                return (slice_p, slice_q,
                        up[index == return_slice_nb],
                        uq[index == return_slice_nb])

    @property
    def local_order(self):
        def alpha(arg1, arg2):
            arg_i = LazyTensor(arg1[:, None, :])
            arg_j = LazyTensor(arg2[None, :, :])
            return ((arg_i * arg_j).sum(-1)) ** 2

        Phi = self.nonlinear_local_average(alpha, self.bo, self.bo)
        Nneigh = self.number_of_neighbours()
        Phi = Phi - (1. / self.N)
        Nneigh = torch.max(Nneigh - (1. / self.N),
                           torch.tensor(1. / self.N,
                                        dtype=Nneigh.dtype,
                                        device=Nneigh.device))
        order_i = Phi * (Nneigh ** (-1))
        return (1 / self.N) * order_i.sum().item()

    @property
    def global_order(self):
        r"""Returns

         .. math::

            \frac{1}{ N(N-1)} \sum_{i,j} (q_i \cdot q_j)^2.

        """
        bo_i = LazyTensor(self.bo[:, None, :])
        bo_j = LazyTensor(self.bo[None, :, :])
        bo_ij = ((bo_i * bo_j).sum(-1)) ** 2
        order_i = 1. / (self.N - 1.) * (bo_ij.sum(1) - 1.)
        return (1. / self.N) * order_i.sum().item()

    @property
    def order_parameter(self):
        return self.global_order

    def build_reflection_quat(self, who, axis):
        r"""

        Returns the quaternion associated to the rotation of angle
        :math:`2\alpha` and axis :math:`n` where :math:`\alpha` is the
        angle between ``vel[who,:]`` and the plane  orthogonal to
        **axis** and :math:`n` is the element of this plane obtained by
        a rotation of :math:`\pi/2` around **axis** of the normalised
        projection of ``vel[who,:]``.

        Args:
            who ((N,) BoolTensor)
            axis (int): The axis number either 0 (:math:`x` axis), 1
                (:math:`y` axis) or 2 (:math:`z` axis).

        Returns:
            (M,4) Tensor:

            Rotation (unit quaternion).

        """
        Omega = self.vel[who, :]
        Om1 = Omega[:, 0]
        Om2 = Omega[:, 1]
        Om3 = Omega[:, 2]
        K = Omega.shape[0]
        if axis == 0:
            q0 = torch.sqrt(Om2 ** 2 + Om3 ** 2)
            q = torch.stack(
                (q0,
                 torch.zeros(K, dtype=self.bo.dtype, device=self.bo.device),
                 -Om1 * Om3 / q0,
                 Om1 * Om2 / q0), dim=1)
        elif axis == 1:
            q0 = torch.sqrt(Om3 ** 2 + Om1 ** 2)
            q = torch.stack(
                (q0,
                 -Om2 * Om3 / q0,
                 torch.zeros(K, dtype=self.bo.dtype, device=self.bo.device),
                 Om2 * Om1 / q0), dim=1)
        elif axis == 2:
            q0 = torch.sqrt(Om1 ** 2 + Om2 ** 2)
            q = torch.stack(
                (q0,
                 -Om3 * Om2 / q0,
                 Om3 * Om1 / q0,
                 torch.zeros(K, dtype=self.bo.dtype, device=self.bo.device)),
                dim=1)
        else:
            raise ValueError('axis must be 0, 1 or 2.')
        return q

    def check_boundary(self):
        r"""Update the positions and body-orientations of the particles
           which are outside the boundary.

        """
        if self.bc == 'periodic':
            for i in range(self.d):
                self.pos[:, i] = torch.remainder(self.pos[:, i], self.L[i])
        elif type(self.bc) == list:
            for i in range(self.d):
                if self.bc[i] == 1:
                    inf0 = self.pos[:, i] < 0
                    supL = self.pos[:, i] > self.L[i]
                    self.pos[inf0, i] = torch.max(
                        -self.pos[inf0, i],
                        torch.tensor([0.],
                                     dtype=self.pos.dtype,
                                     device=self.pos.device))
                    self.pos[supL, i] = torch.min(
                        2 * self.L[i] - self.pos[supL, i],
                        self.L[i])

                    q_infzero = self.build_reflection_quat(who=inf0, axis=i)
                    self.bo[inf0, :] = quat_mult(q_infzero,
                                                 self.bo[inf0, :])
                    q_supL = self.build_reflection_quat(who=supL, axis=i)
                    self.bo[supL, :] = quat_mult(q_supL,
                                                 self.bo[supL, :])
                elif self.bc[0] == 0:
                    self.pos[:, i] = torch.remainder(self.pos[:, i], self.L[i])
                else:
                    raise NotImplementedError()
            # Normalise for numerical stability
            self.bo = self.bo / torch.sqrt(
                torch.sum(self.bo ** 2, dim=1)).reshape((self.N, 1))
        else:
            raise NotImplementedError()

max_stars_repo_name: antoinediez/Sisyphe

max_issues_repo_name: antoinediez/Sisyphe

max_forks_repo_name: antoinediez/Sisyphe


text:
import numpy as np
from mlir_graphblas.sparse_utils import MLIRSparseTensor

from typing import Sequence

GRAPHBLAS_PASSES = [
    "--graphblas-lower",
    "--sparsification",
    "--sparse-tensor-conversion",
    "--linalg-bufferize",
    "--func-bufferize",
    "--tensor-bufferize",
    "--tensor-constant-bufferize",
    "--finalizing-bufferize",
    "--convert-linalg-to-loops",
    "--convert-scf-to-std",
    "--convert-std-to-llvm",
]

STANDARD_PASSES = [
    "--sparsification",
    "--sparse-tensor-conversion",
    "--linalg-bufferize",
    "--func-bufferize",
    "--tensor-bufferize",
    "--tensor-constant-bufferize",
    "--finalizing-bufferize",
    "--convert-linalg-to-loops",
    "--convert-scf-to-std",
    "--convert-std-to-llvm",
]

MLIR_TYPE_TO_NP_TYPE = {
    "i8": np.int8,
    "i16": np.int16,
    "i32": np.int32,
    "i64": np.int64,
    # 'f16': np.float16, # 16-bit floats don't seem to be supported in ctypes
    "f32": np.float32,
    "f64": np.float64,
}

def sparsify_array(
    input_array: np.ndarray, sparsity_values: Sequence[bool]
) -> MLIRSparseTensor:
    """Converts a numpy array into a MLIRSparseTensor."""

    indices = np.array(
        list(zip(*input_array.nonzero())),
        dtype=np.uint64,
    )
    values = np.array(
        [input_array[coordinate] for coordinate in map(tuple, indices)],
        dtype=input_array.dtype,
    )
    sizes = np.array(input_array.shape, dtype=np.uint64)
    sparsity = np.array(sparsity_values, dtype=np.bool8)

    sparse_tensor = MLIRSparseTensor(indices, values, sizes, sparsity)
    return sparse_tensor

ext: py

reason: import numpy


text:
# vim: expandtab:ts=4:sw=4
import numpy as np
import scipy.linalg

"""
Table for the 0.95 quantile of the chi-square distribution with N degrees of
freedom (contains values for N=1, ..., 9). Taken from MATLAB/Octave's chi2inv
function and used as Mahalanobis gating threshold.
"""
chi2inv95 = {
    1: 3.8415,
    2: 5.9915,
    3: 7.8147,
    4: 9.4877,
    5: 11.070,
    6: 12.592,
    7: 14.067,
    8: 15.507,
    9: 16.919}

class KalmanFilter(object):
    """
    A simple Kalman filter for tracking bounding boxes in image space.

    The 8-dimensional state space

        x, y, a, h, vx, vy, va, vh

    contains the bounding box center position (x, y), aspect ratio a, height h,
    and their respective velocities.

    Object motion follows a constant velocity model. The bounding box location
    (x, y, a, h) is taken as direct observation of the state space (linear
    observation model).

    """

    def __init__(self):
        ndim, dt = 4, 1.

        # Create Kalman filter model matrices.
        self._motion_mat = np.eye(2 * ndim, 2 * ndim)
        for i in range(ndim):
            self._motion_mat[i, ndim + i] = dt
        self._update_mat = np.eye(ndim, 2 * ndim)

        # Motion and observation uncertainty are chosen relative to the current
        # state estimate. These weights control the amount of uncertainty in
        # the model. This is a bit hacky.
        self._std_weight_position = 1. / 20
        self._std_weight_velocity = 1. / 160

    def init(self, image, roi):
        measurement = np.array(roi)
        measurement[:2] += measurement[2:] / 2
        measurement[2] /= measurement[3]
        mean_pos = measurement
        mean_vel = np.zeros_like(mean_pos)
        mean = np.r_[mean_pos, mean_vel]

        std = [
            2 * self._std_weight_position * measurement[3],
            2 * self._std_weight_position * measurement[3],
            1e-2,
            2 * self._std_weight_position * measurement[3],
            10 * self._std_weight_velocity * measurement[3],
            10 * self._std_weight_velocity * measurement[3],
            1e-5,
            10 * self._std_weight_velocity * measurement[3]]
        covariance = np.diag(np.square(std))

        self.mean, self.covariance = mean, covariance

    def update(self, image, vis=False):
        std_pos = [
            self._std_weight_position * self.mean[3],
            self._std_weight_position * self.mean[3],
            1e-2,
            self._std_weight_position * self.mean[3]]
        std_vel = [
            self._std_weight_velocity * self.mean[3],
            self._std_weight_velocity * self.mean[3],
            1e-5,
            self._std_weight_velocity * self.mean[3]]
        motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))

        mean = np.dot(self._motion_mat, self.mean)
        covariance = np.linalg.multi_dot((
            self._motion_mat, self.covariance, self._motion_mat.T)) + motion_cov
        self.mean = mean
        self.covariance = covariance

        return self.mean_to_bbox()

    def mean_to_bbox(self):
        ret = self.mean[:4].copy()
        ret[2] *= ret[3]
        ret[:2] -= ret[2:] / 2
        return ret

max_stars_repo_name: ElnuraMusaoglu/SingleObjectTracking

max_issues_repo_name: ElnuraMusaoglu/SingleObjectTracking

max_forks_repo_name: ElnuraMusaoglu/SingleObjectTracking

reason: import numpy,import scipy


text:
import numpy as np
import torch
from torch import optim
from torch.nn import functional
import xgboost as xgb
from conditional import ConditionalGenerator
from models import Discriminator, Generator
from sampler import Sampler
from transformer import DataTransformer
import pickle
import pandas as pd

class CTGANSynthesizer(object):
    """Conditional Table GAN Synthesizer.

    This is the core class of the CTGAN project, where the different components
    are orchestrated together.

    For more details about the process, please check the [Modeling Tabular data using
    Conditional GAN](https://arxiv.org/abs/1907.00503) paper.

    Args:
        embedding_dim (int):
            Size of the random sample passed to the Generator. Defaults to 128.
        gen_dim (tuple or list of ints):
            Size of the output samples for each one of the Residuals. A Resiudal Layer
            will be created for each one of the values provided. Defaults to (256, 256).
        dis_dim (tuple or list of ints):
            Size of the output samples for each one of the Discriminator Layers. A Linear Layer
            will be created for each one of the values provided. Defaults to (256, 256).
        l2scale (float):
            Wheight Decay for the Adam Optimizer. Defaults to 1e-6.
        batch_size (int):
            Number of data samples to process in each step.
    """

    def __init__(self, embedding_dim=128, gen_dim=(256, 256), dis_dim=(256, 256),
                 l2scale=1e-6, batch_size=500):

        self.embedding_dim = embedding_dim
        self.gen_dim = gen_dim
        self.dis_dim = dis_dim

        self.l2scale = l2scale
        self.batch_size = batch_size
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    def _apply_activate(self, data):
        data_t = []
        st = 0
        for item in self.transformer.output_info:
            if item[1] == 'tanh':
                ed = st + item[0]
                data_t.append(torch.tanh(data[:, st:ed]))
                st = ed
            elif item[1] == 'softmax':
                ed = st + item[0]
                data_t.append(functional.gumbel_softmax(data[:, st:ed], tau=0.2))
                st = ed
            else:
                assert 0

        return torch.cat(data_t, dim=1)

    def _cond_loss(self, data, c, m):
        loss = []
        st = 0
        st_c = 0
        skip = False
        for item in self.transformer.output_info:
            if item[1] == 'tanh':
                st += item[0]
                skip = True

            elif item[1] == 'softmax':
                if skip:
                    skip = False
                    st += item[0]
                    continue

                ed = st + item[0]
                ed_c = st_c + item[0]
                tmp = functional.cross_entropy(
                    data[:, st:ed],
                    torch.argmax(c[:, st_c:ed_c], dim=1),
                    reduction='none'
                )
                loss.append(tmp)
                st = ed
                st_c = ed_c

            else:
                assert 0

        loss = torch.stack(loss, dim=1)

        return (loss * m).sum() / data.size()[0]

    def fit(self, train_data, prefered_label, black_box_path, discrete_columns=tuple(),conditional_cols = None, epochs=300, log_frequency=True):
        """Fit the CTGAN Synthesizer models to the training data.

        Args:
            train_data (numpy.ndarray or pandas.DataFrame):
                Training Data. It must be a 2-dimensional numpy array or a
                pandas.DataFrame.
            discrete_columns (list-like):
                List of discrete columns to be used to generate the Conditional
                Vector. If ``train_data`` is a Numpy array, this list should
                contain the integer indices of the columns. Otherwise, if it is
                a ``pandas.DataFrame``, this list should contain the column names.
            epochs (int):
                Number of training epochs. Defaults to 300.
            log_frequency (boolean):
                Whether to use log frequency of categorical levels in conditional
                sampling. Defaults to ``True``.
        """

        self.prefered_label = prefered_label
        self.blackbox_model = pickle.load(open(black_box_path, "rb"))

        self.transformer = DataTransformer()
        self.transformer.fit(train_data, discrete_columns)
        train_data = self.transformer.transform(train_data)

        data_sampler = Sampler(train_data, self.transformer.output_info)

        data_dim = self.transformer.output_dimensions
        self.cond_generator = ConditionalGenerator(
            train_data,
            self.transformer.output_info,
            log_frequency, conditional_cols
        )

        self.generator = Generator(
            self.embedding_dim + self.cond_generator.n_opt,
            self.gen_dim,
            data_dim
        ).to(self.device)

        discriminator = Discriminator(
            data_dim,
            self.dis_dim,1
        ).to(self.device)

        conditonal_discriminator = Discriminator(
            1 + self.cond_generator.n_opt,
            self.dis_dim
        ).to(self.device)

        optimizerG = optim.Adam(
            self.generator.parameters(), lr=2e-4, betas=(0.5, 0.9),
            weight_decay=self.l2scale
        )
        optimizerD = optim.Adam(discriminator.parameters(), lr=2e-4, betas=(0.5, 0.9))
        optimizerconditionalD = optim.Adam(conditonal_discriminator.parameters(), lr=2e-4, betas=(0.5, 0.9))

        assert self.batch_size % 2 == 0
        mean = torch.zeros(self.batch_size, self.embedding_dim, device=self.device)
        std = mean + 1

        steps_per_epoch = max(len(train_data) // self.batch_size, 1)
        for i in range(epochs):
            flip_loss_list = []
            real_flip_loss_list = []
            for id_ in range(steps_per_epoch):

                fakez = torch.normal(mean=mean, std=std)
                condvec = self.cond_generator.sample(self.batch_size)
                if condvec is None:
                    c1, m1, col, opt = None, None, None, None
                    real = data_sampler.sample(self.batch_size, col, opt)
                else:
                    c1, m1, col, opt = condvec
                    c1 = torch.from_numpy(c1).to(self.device)
                    m1 = torch.from_numpy(m1).to(self.device)
                    fakez = torch.cat([fakez, c1], dim=1)

                    perm = np.arange(self.batch_size)
                    np.random.shuffle(perm)
                    real = data_sampler.sample(self.batch_size, col[perm], opt[perm])
                    c2 = c1[perm]

                fake = self.generator(fakez)
                fakeact = self._apply_activate(fake)

                if condvec is None:
                    cross_entropy = 0
                else:
                    cross_entropy = self._cond_loss(fake, c1, m1)

                real = torch.from_numpy(real.astype('float32')).to(self.device)

                if c1 is not None:

                    real_cat = real
                    fake_cat = fakeact

                else:
                    real_cat = real
                    fake_cat = fake

                y_fake = discriminator(fake_cat)
                y_real = discriminator(real_cat)

                if c1 is not None:
                    conditional_fake_cat = torch.cat([y_fake, c1], dim=1)
                    conditional_real_cat = torch.cat([y_real, c2], dim=1)

                else:
                    conditional_fake_cat = y_fake
                    conditional_real_cat = y_real

                conditional_y_fake = conditonal_discriminator(conditional_fake_cat)
                conditional_y_real = conditonal_discriminator(conditional_real_cat)

                pen = discriminator.calc_gradient_penalty(real_cat, fake_cat, self.device)
                loss_d = -(torch.mean(y_real) - torch.mean(y_fake))

                condtional_pen = conditonal_discriminator.calc_gradient_penalty(conditional_real_cat, conditional_fake_cat, self.device)
                loss_condtional_d = -(torch.mean(conditional_y_real) - torch.mean(conditional_y_fake))

                optimizerD.zero_grad()
                pen.backward(retain_graph=True)
                loss_d.backward(retain_graph=True)
                optimizerD.step()

                optimizerconditionalD.zero_grad()
                condtional_pen.backward(retain_graph=True)
                loss_condtional_d.backward()
                optimizerconditionalD.step()

                fakez = torch.normal(mean=mean, std=std)
                condvec = self.cond_generator.sample(self.batch_size)

                if condvec is None:
                    c1, m1, col, opt = None, None, None, None
                else:
                    c1, m1, col, opt = condvec
                    c1 = torch.from_numpy(c1).to(self.device)
                    m1 = torch.from_numpy(m1).to(self.device)
                    fakez = torch.cat([fakez, c1], dim=1)

                fake = self.generator(fakez)
                fakeact = self._apply_activate(fake)

                if c1 is not None:
                    y_fake = discriminator(fakeact)
                    conditional_y_fake = conditonal_discriminator(torch.cat([y_fake, c1], dim=1))
                else:
                    y_fake = discriminator(fakeact)
                    conditional_y_fake = conditonal_discriminator(y_fake)

                if condvec is None:
                    cross_entropy = 0
                else:
                    cross_entropy = self._cond_loss(fake, c1, m1)

                fake_act_inv = self.transformer.inverse_transform(fakeact.detach().cpu().numpy(), None)
                fake_act_inv = _factorize_categoricals(fake_act_inv, discrete_columns)
                fake_act_inv = xgb.DMatrix(data = fake_act_inv)
                black_box_pred_prob = self.blackbox_model.predict(fake_act_inv)
                #black_box_pred_prob = torch.from_numpy(np.stack([1-black_box_pred_prob,black_box_pred_prob], axis = -1))
                #flip_loss = torch.nn.CrossEntropyLoss()(black_box_pred_prob, torch.tensor([self.prefered_label]).repeat(self.batch_size))
                flip_loss = sum(-np.log(black_box_pred_prob))/self.batch_size

                real_inv = self.transformer.inverse_transform(real.detach().cpu().numpy(), None)
                real_inv = _factorize_categoricals(real_inv, discrete_columns)
                real_inv = xgb.DMatrix(data = real_inv)
                real_pred_prob = self.blackbox_model.predict(real_inv)
                #black_box_pred_prob = torch.from_numpy(np.stack([1-black_box_pred_prob,black_box_pred_prob], axis = -1))
                #flip_loss = torch.nn.CrossEntropyLoss()(black_box_pred_prob, torch.tensor([self.prefered_label]).repeat(self.batch_size))
                real_flip_loss = sum(-np.log(real_pred_prob))/self.batch_size

                loss_g = -torch.mean(conditional_y_fake) +  cross_entropy + 10*flip_loss
                #print(f"Base Loss:{-torch.mean(conditional_y_fake)}, Conditional Loss:{cross_entropy}, 10Flip Loss:{flip_loss}")
                flip_loss_list.append(flip_loss)
                real_flip_loss_list.append(real_flip_loss)

                optimizerG.zero_grad()
                loss_g.backward()
                optimizerG.step()

            print(f"Generated flip loss {np.mean(flip_loss_list)}, Real flip loss {np.mean(real_flip_loss_list)}")
            print("Condtional Cross Entropy Loss", cross_entropy)
            print("Epoch %d, Loss G: %.4f, Loss D: %.4f, Loss Conditional D: %.4f" %
                  (i + 1, loss_g.detach().cpu(), loss_d.detach().cpu(), loss_condtional_d.detach().cpu()),
                  flush=True)

    def sample(self, n, col_index = None):
        """Sample data similar to the training data.
        Args:
            n (int):
                Number of rows to sample.
        Returns:
            numpy.ndarray or pandas.DataFrame
        """

        steps = n // self.batch_size + 1
        data = []
        for i in range(steps):
            mean = torch.zeros(self.batch_size, self.embedding_dim)
            std = mean + 1
            fakez = torch.normal(mean=mean, std=std).to(self.device)

            condvec,m1 = self.cond_generator.sample_zero(self.batch_size)
            m1 = torch.from_numpy(m1).to(self.device)
            if condvec is None:
                pass
            else:
                c1 = condvec
                if col_index != None:
                    c1 = np.zeros_like(c1)
                    c1[:,col_index] = 1
                c1 = torch.from_numpy(c1).to(self.device)
                fakez = torch.cat([fakez, c1], dim=1)

            fake = self.generator(fakez)
            fakeact = self._apply_activate(fake)
            data.append(fakeact.detach().cpu().numpy())
            print(self._cond_loss(fake, c1, m1))

        data = np.concatenate(data, axis=0)
        data = data[:n]

        return self.transformer.inverse_transform(data, None)

def _factorize_categoricals(df, discrete_columns):
    for col in discrete_columns:
        df[col], _ = pd.factorize(df[col])
    return df

max_stars_repo_name: sahanalva/CTGAN

max_issues_repo_name: sahanalva/CTGAN

max_forks_repo_name: sahanalva/CTGAN


text:
#!/usr/bin/env python
u"""
geocenter_compare_tellus.py
Written by Tyler Sutterley (12/2021)
Plots the GRACE/GRACE-FO geocenter time series for different
    GRACE/GRACE-FO processing centers comparing with the
    JPL GRACE Tellus product

CALLING SEQUENCE:
    python geocenter_processing_centers.py --start 4 --end 216

COMMAND LINE OPTIONS:
    -D X, --directory X: working data directory with geocenter files
    -r X, --release X: GRACE/GRACE-FO data release
    -S X, --start X: starting GRACE month for time series
    -E X, --end X: ending GRACE month for time series
    -M X, --missing X: Missing GRACE months in time series

UPDATE HISTORY:
    Updated 12/2021: adjust minimum x limit based on starting GRACE month
    Updated 11/2021: use gravity_toolkit geocenter class for operations
    Written 05/2021
"""
from __future__ import print_function

import os
import argparse
import numpy as np
import matplotlib
import matplotlib.font_manager
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, FormatStrFormatter
from matplotlib.offsetbox import AnchoredText,AnchoredOffsetbox,TextArea,VPacker
from gravity_toolkit.time import convert_calendar_decimal
import gravity_toolkit.geocenter as geocenter

#-- rebuilt the matplotlib fonts and set parameters
matplotlib.font_manager._load_fontmanager()
matplotlib.rcParams['font.family'] = 'sans-serif'
matplotlib.rcParams['font.sans-serif'] = ['Helvetica']
matplotlib.rcParams['mathtext.default'] = 'regular'

#-- PURPOSE: plots the GRACE/GRACE-FO geocenter time series
def geocenter_compare_tellus(grace_dir,DREL,START_MON,END_MON,MISSING):
    #-- GRACE months
    GAP = [187,188,189,190,191,192,193,194,195,196,197]
    months = sorted(set(np.arange(START_MON,END_MON+1)) - set(MISSING))
    #-- labels for each scenario
    input_flags = ['','iter','SLF_iter','SLF_iter_wSLR21']
    input_labels = ['Static','Iterated','Iterated SLF']
    #-- labels for Release-6
    PROC = ['CSR','GFZ','JPL']
    model_str = 'OMCT' if DREL in ('RL04','RL05') else 'MPIOM'
    #-- degree one coefficient labels
    fig_labels = ['C11','S11','C10']
    axes_labels = dict(C10='c)',C11='a)',S11='b)')
    ylabels = dict(C10='z',C11='x',S11='y')

    #-- plot colors for each dataset
    plot_colors = {'Iterated SLF':'darkorchid','GFZ GravIS':'darkorange',
        'JPL Tellus':'mediumseagreen'}

    #-- plot geocenter estimates for each processing center
    for k,pr in enumerate(PROC):
        #-- 3 row plot (C10, C11 and S11)
        ax = {}
        fig,(ax[0],ax[1],ax[2])=plt.subplots(num=1,ncols=3,
            sharey=True,figsize=(9,4))
        #-- additionally plot GFZ with SLR replaced pole tide
        if (pr == 'GFZwPT'):
            fargs = ('GFZ',DREL,model_str,input_flags[3])
        else:
            fargs = (pr,DREL,model_str,input_flags[2])
        #-- read geocenter file for processing center and model
        grace_file = '{0}_{1}_{2}_{3}.txt'.format(*fargs)
        DEG1 = geocenter().from_UCI(os.path.join(grace_dir,grace_file))
        #-- indices for mean months
        kk, = np.nonzero((DEG1.month >= START_MON) & (DEG1.month <= 176))
        DEG1.mean(apply=True, indices=kk)
        #-- setting Load Love Number (kl) to 0.021 to match Swenson et al. (2008)
        DEG1.to_cartesian(kl=0.021)
        #-- plot each coefficient
        for j,key in enumerate(fig_labels):
            #-- plot model outputs
            #-- create a time series with nans for missing months
            tdec = np.full_like(months,np.nan,dtype=np.float64)
            data = np.full_like(months,np.nan,dtype=np.float64)
            val = getattr(DEG1, ylabels[key].upper())
            for i,m in enumerate(months):
                valid = np.count_nonzero(DEG1.month == m)
                if valid:
                    mm, = np.nonzero(DEG1.month == m)
                    tdec[i] = DEG1.time[mm]
                    data[i] = val[mm]
            #-- plot all dates
            ax[j].plot(tdec, data, color=plot_colors['Iterated SLF'],
                label='Iterated SLF')

        if (pr == 'GFZwPT'):
            grace_file = 'GRAVIS-2B_GFZOP_GEOCENTER_0002.dat'
            DEG1 = geocenter().from_gravis(os.path.join(grace_dir,grace_file))
            #-- indices for mean months
            kk, = np.nonzero((DEG1.month >= START_MON) & (DEG1.month <= 176))
            DEG1.mean(apply=True, indices=kk)
            #-- setting Load Love Number (kl) to 0.021 to match Swenson et al. (2008)
            DEG1.to_cartesian(kl=0.021)
            #-- plot each coefficient
            for j,key in enumerate(fig_labels):
                #-- plot model outputs
                DEG1[key] -= DEG1[key][kk].mean()
                #-- create a time series with nans for missing months
                tdec = np.full_like(months,np.nan,dtype=np.float64)
                data = np.full_like(months,np.nan,dtype=np.float64)
                val = getattr(DEG1, ylabels[key].upper())
                for i,m in enumerate(months):
                    valid = np.count_nonzero(DEG1.month == m)
                    if valid:
                        mm, = np.nonzero(DEG1.month == m)
                        tdec[i] = DEG1.time[mm]
                        data[i] = val[mm]
                #-- plot all dates
                ax[j].plot(tdec, data, color=plot_colors['GFZ GravIS'],
                    label='GFZ GravIS')

        #-- Running function read_tellus_geocenter.py
        grace_file = 'TN-13_GEOC_{0}_{1}.txt'.format(pr,DREL)
        DEG1 = geocenter().from_tellus(os.path.join(grace_dir,grace_file),JPL=True)
        #-- indices for mean months
        kk, = np.nonzero((DEG1.month >= START_MON) & (DEG1.month <= 176))
        DEG1.mean(apply=True, indices=kk)
        #-- setting Load Love Number (kl) to 0.021 to match Swenson et al. (2008)
        DEG1.to_cartesian(kl=0.021)
        #-- plot each coefficient
        for j,key in enumerate(fig_labels):
            #-- plot model outputs
            #-- create a time series with nans for missing months
            tdec = np.full_like(months,np.nan,dtype=np.float64)
            data = np.full_like(months,np.nan,dtype=np.float64)
            val = getattr(DEG1, ylabels[key].upper())
            for i,m in enumerate(months):
                valid = np.count_nonzero(DEG1.month == m)
                if valid:
                    mm, = np.nonzero(DEG1.month == m)
                    tdec[i] = DEG1.time[mm]
                    data[i] = val[mm]
            #-- plot all dates
            ax[j].plot(tdec, data, color=plot_colors['JPL Tellus'],
                label='JPL Tellus')

        #-- add axis labels and adjust font sizes for axis ticks
        for j,key in enumerate(fig_labels):
            #-- vertical line denoting the accelerometer shutoff
            acc = convert_calendar_decimal(2016,9,day=3,hour=12,minute=12)
            ax[j].axvline(acc,color='0.5',ls='dashed',lw=0.5,dashes=(8,4))
            #-- vertical lines for end of the GRACE mission and start of GRACE-FO
            jj, = np.flatnonzero(DEG1.month == 186)
            kk, = np.flatnonzero(DEG1.month == 198)
            vs = ax[j].axvspan(DEG1.time[jj],DEG1.time[kk],
                color='0.5',ls='dashed',alpha=0.15)
            vs._dashes = (4,2)
            #-- axis label
            ax[j].set_title(ylabels[key], style='italic', fontsize=14)
            ax[j].add_artist(AnchoredText(axes_labels[key], pad=0.,
                prop=dict(size=16,weight='bold'), frameon=False, loc=2))
            ax[j].set_xlabel('Time [Yr]', fontsize=14)
            #-- set ticks
            xmin = 2002 + (START_MON + 1.0)//12.0
            xmax = 2002 + (END_MON + 1.0)/12.0
            major_ticks = np.arange(2005, xmax, 5)
            ax[j].xaxis.set_ticks(major_ticks)
            minor_ticks = sorted(set(np.arange(xmin, xmax, 1)) - set(major_ticks))
            ax[j].xaxis.set_ticks(minor_ticks, minor=True)
            ax[j].set_xlim(xmin, xmax)
            ax[j].set_ylim(-9.5,8.5)
            #-- axes tick adjustments
            ax[j].get_xaxis().set_tick_params(which='both', direction='in')
            ax[j].get_yaxis().set_tick_params(which='both', direction='in')
            for tick in ax[j].xaxis.get_major_ticks():
                tick.label.set_fontsize(14)
            for tick in ax[j].yaxis.get_major_ticks():
                tick.label.set_fontsize(14)

        #-- add legend
        lgd = ax[0].legend(loc=3,frameon=False)
        lgd.get_frame().set_alpha(1.0)
        for line in lgd.get_lines():
            line.set_linewidth(6)
        for i,text in enumerate(lgd.get_texts()):
            text.set_weight('bold')
            text.set_color(plot_colors[text.get_text()])
        #-- labels and set limits
        ax[0].set_ylabel('Geocenter Variation [mm]', fontsize=14)
        #-- adjust locations of subplots
        fig.subplots_adjust(left=0.06,right=0.98,bottom=0.12,top=0.94,wspace=0.05)
        #-- save figure to file
        OUTPUT_FIGURE = 'TN13_SV19_{0}_{1}.pdf'.format(pr,DREL)
        plt.savefig(os.path.join(grace_dir,OUTPUT_FIGURE), format='pdf', dpi=300)
        plt.clf()

#-- This is the main part of the program that calls the individual modules
def main():
    #-- Read the system arguments listed after the program
    parser = argparse.ArgumentParser(
        description="""Plots the GRACE/GRACE-FO geocenter time series for
            different GRACE/GRACE-FO processing centers comparing with the
            JPL GRACE Tellus product
            """
    )
    #-- working data directory
    parser.add_argument('--directory','-D',
        type=lambda p: os.path.abspath(os.path.expanduser(p)),
        default=os.getcwd(),
        help='Working data directory')
    #-- GRACE/GRACE-FO data release
    parser.add_argument('--release','-r',
        metavar='DREL', type=str,
        default='RL06', choices=['RL04','RL05','RL06'],
        help='GRACE/GRACE-FO data release')
    #-- start and end GRACE/GRACE-FO months
    parser.add_argument('--start','-S',
        type=int, default=4,
        help='Starting GRACE/GRACE-FO month for time series')
    parser.add_argument('--end','-E',
        type=int, default=231,
        help='Ending GRACE/GRACE-FO month for time series')
    MISSING = [6,7,18,109,114,125,130,135,140,141,146,151,156,162,166,167,172,
        177,178,182,200,201]
    parser.add_argument('--missing','-M',
        metavar='MISSING', type=int, nargs='+', default=MISSING,
        help='Missing GRACE/GRACE-FO months in time series')
    args,_ = parser.parse_known_args()

    #-- run program with parameters
    geocenter_compare_tellus(args.directory, args.release,
        args.start, args.end, args.missing)

#-- run main program
if __name__ == '__main__':
    main()

max_stars_repo_name: richannan/read-GRACE-harmonics

max_issues_repo_name: richannan/read-GRACE-harmonics

max_forks_repo_name: richannan/read-GRACE-harmonics


text:
import numpy as np

# ----------
# Particle Class
class Particle(object):
    def __init__(self):
        self.x = np.random.random((1,2)) # [[x,y]] coordinates
        self.xp = np.zeros((1,2))
        self.radius = 0.1
        self.mass = 1

    def accelerate(self, force_vector, step_size):
        acceleration = np.multiply(force_vector, (1/self.mass))
        self.xp = np.add(self.xp, np.multiply(step_size, acceleration))
        return self

    def move(self, step_size):
        self.x = np.add(self.x, np.multiply(step_size, self.xp))
        return self

    def edge_check(self, xmax, ymax):
        self.x = self.x % 1
        return self

    def set_velocity(self, velocity):
        self.xp = velocity
        return self

    def set_position(self, position):
        self.x = position
        return self

max_stars_repo_name: ltwp/sandstorm

max_issues_repo_name: ltwp/sandstorm

max_forks_repo_name: ltwp/sandstorm