text:
"""
Module for handling molecules. Uses the pymatgen.core.structure.Molecule
class as a base. Has a function for reorienting molecules
(reoriented_molecule), and for calculating valid orientations within a Wyckoff
position based on symmetry (orientation_in_wyckoff_position).
The orientation class can be used to identify
degrees of freedom for molecules in Wyckoff positions with certain symmetry
constraints.

"""
# Imports
import numpy as np
from copy import deepcopy
from scipy.spatial.transform import Rotation
import networkx as nx
from random import choice
# ------------------------------
# External Libraries
from pymatgen.core.structure import Molecule
from pymatgen.symmetry.analyzer import PointGroupAnalyzer, generate_full_symmops
from pymatgen.core.bonds import CovalentBond

# PyXtal imports
from pyxtal.msg import printx
from pyxtal.tolerance import Tol_matrix
from pyxtal.database.element import Element
from pyxtal.operations import SymmOp, OperationAnalyzer, rotate_vector, angle
from pyxtal.database.collection import Collection

# Define functions
# ------------------------------
molecule_collection = Collection("molecules")

class pyxtal_molecule:
    """
    Extended molecule class based on pymatgen.core.structure.Molecule
    The added features include:
    0, parse the input
    1, estimate volume/tolerance/radii
    2, find and store symmetry (todo)
    3, get the principle axis (todo)
    4, re-align the molecule (todo)

    Args:
        mol: a string to reprent the molecule
        tm: tolerance matrix

    """

    def __init__(self, mol=None, symmetrize=True, tm=Tol_matrix(prototype="molecular")):
        mo = None
        if type(mol) == str:
            # Parse molecules: either file or molecule name
            tmp = mol.split(".")
            self.name = tmp[0]
            if len(tmp) > 1:
                # Load the molecule from the given file
                if tmp[-1] in ["xyz", "gjf", "g03", "json"]:
                    import os
                    if os.path.exists(mol):
                        mo = Molecule.from_file(mol)
                    else:
                        raise NameError("{:s} is not a valid path".format(mol))
                else:
                    raise NameError("{:s} is not a supported format".format(tmp[-1]))
            else:
                # print('\nLoad the molecule {:s} from collections'.format(mol))
                mo = molecule_collection[mol]
        elif hasattr(mol, "sites"):  # pymatgen molecule
            self.name = str(mol.formula)
            mo = mol

        if mo is None:
            msg = "Could not create molecules from given input: {:s}".format(mol)
            raise NameError(msg)

        self.props = mo.site_properties

        if len(mo) > 1:
            if symmetrize:
                pga = PointGroupAnalyzer(mo)
                mo = pga.symmetrize_molecule()["sym_mol"]
            mo = self.add_site_props(mo)

        self.mol = mo
        self.tm = tm
        self.get_box()
        self.volume = self.box.volume
        self.get_radius()
        self.get_symbols()
        self.get_tols_matrix()

    def __str__(self):
        return '[' + self.name + ']'

    def save_dict(self):
        return self.mol.as_dict()

    def copy(self):
        """
        simply copy the structure
        """
        return deepcopy(self)

    def reset_positions(self, coors):
        """
        reset the coordinates
        """
        from pymatgen.core.sites import Site
        if len(coors) != len(self.mol._sites):
            raise ValueError("number of atoms is inconsistent!")
        else:
            for i, coor in enumerate(coors):
                _site = self.mol._sites[i]
                new_site = Site(_site.species, coor, properties=_site.properties)
                self.mol._sites[i] = new_site

    @classmethod
    def load_dict(cls, dicts):
        """
        load the molecule from a dictionary
        """
        mol = Molecule.from_dict(dicts)
        return cls(mol)

    def swap_axis(self, ax):
        """
        swap the molecular axis
        """
        coords = self.mol.cart_coords[:, ax]
        mo = Molecule(self.symbols, coords)
        mo = self.add_site_props(mo)

        return pyxtal_molecule(mo, self.tm)

    def add_site_props(self, mo):
        if len(self.props) > 0:
            for key in self.props.keys():
                mo.add_site_property(key, self.props[key])
        return mo

    def get_box(self):
        """
        Given a molecule, find a minimum orthorhombic box containing it.
        Size is calculated using min and max x, y, and z values,
        plus the padding defined by the vdw radius
        For best results, call oriented_molecule first.

        Args:
            mol: a pymatgen Molecule object. Should be oriented along its principle axes.

        Returns:
            a Box object
        """
        mol, P = reoriented_molecule(self.mol)
        minx, miny, minz, maxx, maxy, maxz = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
        for p in mol:
            x, y, z = p.coords
            r = Element(p.species_string).vdw_radius
            if x - r < minx:
                minx = x - r
            if y - r < miny:
                miny = y - r
            if z - r < minz:
                minz = z - r
            if x + r > maxx:
                maxx = x + r
            if y + r > maxy:
                maxy = y + r
            if z + r > maxz:
                maxz = z + r
        self.box = Box(minx, maxx, miny, maxy, minz, maxz)
        self.axes = P

    def get_radius(self):
        r_max = 0
        for coord, number in zip(self.mol.cart_coords, self.mol.atomic_numbers):
            radius = (
                np.sqrt(np.sum(coord * coord)) + self.tm.get_tol(number, number) * 0.5
            )
            if radius > r_max:
                r_max = radius
        self.radius = r_max
        # reestimate the radius if it has stick shape
        rmax = max([self.box.width,self.box.height,self.box.length])
        rmin = min([self.box.width,self.box.height,self.box.length])
        if rmax/rmin > 3 and rmax >12:
            self.radius = rmin

    def has_stick_shape(self):
        sizes = [self.box.width,self.box.height,self.box.length]
        sizes.sort()
        if sizes[2]>15: #and sizes[2]/sizes[0]>2 and sizes[2]/sizes[1]>2:
            return True
        else:
            return False

    def get_symbols(self):
        self.symbols = [specie.name for specie in self.mol.species]

    def get_tols_matrix(self):
        """
        Returns: a 2D matrix which is used internally for distance checking.
        """
        numbers = self.mol.atomic_numbers
        tols = np.zeros((len(numbers), len(numbers)))
        for i1, number1 in enumerate(numbers):
            for i2, number2 in enumerate(numbers):
                tols[i1][i2] = self.tm.get_tol(number1, number2)
        if len(self.mol)==1:
            tols *= 0.8 # if only one atom, reduce the tolerance
        self.tols_matrix = tols

    def show(self):
        from pyxtal.viz import display_molecules
        return display_molecules([self.mol])

class Box:
    """
    Class for storing the binding box for a molecule. Box is oriented along the x, y, and
    z axes.

    Args:
        minx: the minimum x value
        maxx: the maximum x value
        miny: the minimum y value
        maxy: the maximum y value
        minz: the minimum z value
        maxz: the maximum z value
    """

    def __init__(self, minx, maxx, miny, maxy, minz, maxz):
        self.minx = float(minx)
        self.maxx = float(maxx)
        self.miny = float(miny)
        self.maxy = float(maxy)
        self.minz = float(minz)
        self.maxz = float(maxz)

        self.width = float(abs(maxx - minx))
        self.length = float(abs(maxy - miny))
        self.height = float(abs(maxz - minz))

        self.minl = min(self.width, self.length, self.height)
        self.maxl = max(self.width, self.length, self.height)
        for x in (self.width, self.length, self.height):
            if x <= self.maxl and x >= self.minl:
                self.midl = x

        self.volume = float(self.width * self.length * self.height)

class Orientation:
    """
    Stores orientations for molecules based on vector constraints.
    Can be stored to regenerate orientations consistent with a given constraint
    vector, without re-calling orientation_in_wyckoff_position. Allows for
    generating orientations which differ only in their rotation about a given
    axis.

    Args:
        matrix: a 3x3 rotation matrix describing the orientation (and/or
            inversion) to store
        degrees: the number of degrees of freedom...
            0 - The orientation refers to a single rotation matrix
            1 - The orientation can be rotated about a single axis
            2 - The orientation can be any pure rotation matrix

        axis:
            an optional axis about which the orientation can rotate. Only used
            if degrees is equal to 1
    """

    def __init__(self, matrix=None, degrees=2, axis=None):
        self.matrix = np.array(matrix)
        self.degrees = degrees  # The number of degrees of freedom.
        if degrees == 1:
            if axis is None:
                raise ValueError("axis is required for orientation")
            else:
                axis /= np.linalg.norm(axis)
        self.axis = axis

        self.r = Rotation.from_matrix(self.matrix)  # scipy transform.Rotation class
        self.angle = None

    def __str__(self):
        s = "-------PyXtal.molecule.Orientation class----\n"
        s += "degree of freedom: {:d}\n".format(self.degrees)
        s += "Rotation matrix:\n"
        s += "{:6.3f} {:6.3f} {:6.3f}\n".format(*self.matrix[:,0])
        s += "{:6.3f} {:6.3f} {:6.3f}\n".format(*self.matrix[:,1])
        s += "{:6.3f} {:6.3f} {:6.3f}\n".format(*self.matrix[:,2])
        if self.axis is not None:
            s += "Rotation axis\n"
            s += "{:6.2f} {:6.2f} {:6.3f}\n".format(*self.axis)
        return s

    def reset_matrix(self, matrix):
        self.matrix = matrix
        self.r = Rotation.from_matrix(self.matrix)

    def __repr__(self):
        return str(self)

    def copy(self):
        return deepcopy(self)

    def save_dict(self):
        dict0 = {"matrix": self.matrix,
                 "degrees": self.degrees,
                 "axis": self.axis
                }
        return dict0

    @classmethod
    def load_dict(cls, dicts):
        matrix = dicts['matrix']
        degrees = dicts['degrees']
        axis = dicts['axis']
        return cls(matrix, degrees, axis)

    def change_orientation(self, angle="random", flip=False):
        """
        Allows for specification of an angle (possibly random) to
        rotate about the constraint axis.

        Args:
            angle: an angle to rotate about the constraint axis.
            If "random", chooses a random rotation angle.
            If self.degrees==2, chooses a random rotation matrix.
            If self.degrees==1, only apply on angle
            If self.degrees==0, no change

        """
        if self.degrees >= 1:
            # choose the axis
            if self.axis is None:
                axis = np.random.RandomState().rand(3) - 0.5
                self.axis = axis / np.linalg.norm(axis)

            # parse the angle
            if angle == "random":
                angle = np.random.RandomState().rand() * np.pi * 2
            self.angle = angle

            # update the matrix
            r1 = Rotation.from_rotvec(self.angle * self.axis)

            if self.degrees == 2 and flip:
                if np.random.random()>0.5:
                    ax = choice(['x','y','z'])
                    angle0 = choice([90, 180, 270])
                    r2 = Rotation.from_euler(ax, angle0, degrees=True)
                    r1 = r2*r1
            self.r = r1 * self.r
            #self.r *= r1
            self.matrix = self.r.as_matrix()

    def rotate_by_matrix(self, matrix, ignore_constraint=True):
        """
        rotate

        Args:
            matrix: 3*3 rotation matrix

        """
        if not ignore_constraint:
            if self.degrees == 0:
                raise ValueError("cannot rotate")
            elif self.degrees == 1:
                axis = self.axis
                vec = Rotation.from_matrix(matrix).as_rotvec()
                if angle(vec, self.axis) > 1e-2 and angle(vec, -self.axis) > 1e-2:
                    raise ValueError("must rotate along the given axis")
        else:
            axis = None

        matrix = matrix.dot(self.matrix)
        return Orientation(matrix, self.degrees, axis)

    def get_matrix(self, angle="random"):
        """
        Generate a 3x3 rotation matrix consistent with the orientation's
        constraints. Allows for specification of an angle (possibly random) to
        rotate about the constraint axis.

        Args:
            angle: an angle to rotate about the constraint axis. If "random",
                chooses a random rotation angle. If self.degrees==2, chooses a
                random 3d rotation matrix to multiply by. If the original matrix
                is wanted, set angle=0, or call self.matrix

        Returns:
            a 3x3 rotation (and/or inversion) matrix (numpy array)
        """
        if self.degrees == 2:
            if angle == "random":
                axis = np.random.sample(3)
                axis = axis / np.linalg.norm(axis)
                angle = np.random.random() * np.pi * 2
            else:
                axis = self.axis
            return Rotation.from_rotvec(angle * axis).as_matrix()

        elif self.degrees == 1:
            if angle == "random":
                angle = np.random.random() * np.pi * 2
            else:
                angle = self.angle
            return Rotation.from_rotvec(angle * self.axis).as_matrix()

        elif self.degrees == 0:
            return self.matrix

    def get_op(self, angle=None):
        """
        Generate a SymmOp object consistent with the orientation's
        constraints. Allows for specification of an angle (possibly random) to
        rotate about the constraint axis.

        Args:
            angle: an angle to rotate about the constraint axis. If "random",
                chooses a random rotation angle. If self.degrees==2, chooses a
                random 3d rotation matrix to multiply by. If the original matrix
                is wanted, set angle=0, or call self.matrix

        Returns:
            pymatgen.core.structure. SymmOp object
        """
        #if angle is not None:
        #    self.change_orientation(angle)
        return SymmOp.from_rotation_and_translation(self.matrix, [0, 0, 0])

    @classmethod
    def from_constraint(self, v1, c1):
        """
        Geneate an orientation object given a constraint axis c1, and a
        corresponding vector v1. v1 will be rotated onto c1, and the resulting
        orientation will have a rotational degree of freedom about c1.

        Args:
            v1: a 1x3 vector in the original reference frame
            c1: a corresponding axis which v1 must be mapped to

        Returns:
            an orientation object consistent with the supplied constraint
        """
        # c1 is the constraint vector; v1 will be rotated onto it
        m = rotate_vector(v1, c1)
        return Orientation(m, degrees=1, axis=c1)

    @classmethod
    def from_constraints(self, v1, c1, v2, c2):
        """
        Geneate an orientation object given two constraint vectors

        Args:
            v1: a 1x3 vector in the original reference frame
            c1: a corresponding axis which v1 must be mapped to
            v1: a second 1x3 vector in the original reference frame
            c1: a corresponding axis which v2 must be mapped to

        Returns:
            an orientation object consistent with the supplied constraints
        """
        T = rotate_vector(v1, c1)
        phi = angle(c1, c2)
        phi2 = angle(c1, (np.dot(T, v2)))
        if not np.isclose(phi, phi2, rtol=0.01):
            printx("Error: constraints and vectors do not match.", priority=1)
            return
        r = np.sin(phi)
        c = np.linalg.norm(np.dot(T, v2) - c2)
        theta = np.arccos(1 - (c ** 2) / (2 * (r ** 2)))
        Rot = R.from_rotvec(theta * c1)
        T2 = np.dot(R, T)
        a = angle(np.dot(T2, v2), c2)
        if not np.isclose(a, 0, rtol=0.01):
            T2 = np.dot(np.linalg.inv(R), T)
        a = angle(np.dot(T2, v2), c2)
        if not np.isclose(a, 0, rtol=0.01):
            printx("Error: Generated incorrect rotation: " + str(theta), priority=1)
        return Orientation(T2, degrees=0)

    def random_orientation(self):
        """
        Applies random rotation (if possible) and returns a new orientation with
        the new base matrix.

        Returns:
            a new orientation object with a different base rotation matrix
        """

        self.change_orientation()
        return self

    def get_Euler_angles(self):
        return self.r.as_euler('zxy', degrees=True)

def get_inertia_tensor(coords):
    """
    Calculate the symmetric inertia tensor for a Molecule
    the principal axes of symmetry.

    Args:
        coords: [N, 3] array of coordinates

    Returns:
        a 3x3 numpy array representing the inertia tensor
    """
    coords -= np.mean(coords, axis=0)
    Inertia = np.zeros([3,3])
    Inertia[0,0] = np.sum(coords[:,1]**2 + coords[:,2]**2)
    Inertia[1,1] = np.sum(coords[:,0]**2 + coords[:,2]**2)
    Inertia[2,2] = np.sum(coords[:,0]**2 + coords[:,1]**2)
    Inertia[0,1] = Inertia[1,0] = -np.sum(coords[:,0]*coords[:,1])
    Inertia[0,2] = Inertia[2,0] = -np.sum(coords[:,0]*coords[:,2])
    Inertia[1,2] = Inertia[2,1] = -np.sum(coords[:,1]*coords[:,2])

    return Inertia

def reoriented_molecule(mol, nested=False):
    """
    Reorient a molecule so that its principal axes are aligned with the
    identity matrix.

    Args:
        mol: a Molecule object
        nested: keep track of how many times the function
            has been called recursively

    Returns:
        new_mol: a reoriented copy of the original molecule.
        P: the 3x3 rotation matrix used to obtain it.
    """
    coords = mol.cart_coords
    numbers = mol.atomic_numbers
    coords -= np.mean(coords, axis=0)
    A = get_inertia_tensor(coords)
    # Store the eigenvectors of the inertia tensor
    P = np.linalg.eigh(A)[1]
    if np.linalg.det(P) < 0:
        P[0] *= -1
    coords = np.dot(coords, P)
    return Molecule(numbers, coords), P

def get_symmetry(mol, already_oriented=False):
    """
    Return a molecule's point symmetry.
    Note: for linear molecules, infinitessimal rotations are treated as 6-fold
    rotations, which works for 3d and 2d point groups.

    Args:
        mol: a Molecule object
        already_oriented: whether or not the principle axes of mol are already
            reoriented. Can save time if True, but is not required.

    Returns:
        a list of SymmOp objects which leave the molecule unchanged when applied
    """
    # For single atoms, we cannot represent the point group using a list of operations
    if len(mol) == 1:
        return []
    pga = PointGroupAnalyzer(mol)
    # Handle linear molecules
    if "*" in pga.sch_symbol:
        if already_oriented == False:
            # Reorient the molecule
            oriented_mol, P = reoriented_molecule(mol)
            pga = PointGroupAnalyzer(oriented_mol)
        pg = pga.get_pointgroup()
        symm_m = []
        for op in pg:
            symm_m.append(op)
        # Add 12-fold  and reflections in place of ininitesimal rotation
        for axis in [[1, 0, 0], [0, 1, 0], [0, 0, 1]]:
            # op = SymmOp.from_rotation_and_translation(aa2matrix(axis, np.pi/6), [0,0,0])
            m1 = Rotation.from_rotvec(np.pi / 6 * axis).as_matrix()
            op = SymmOp.from_rotation_and_translation(m1, [0, 0, 0])
            if pga.is_valid_op(op):
                symm_m.append(op)
                # Any molecule with infinitesimal symmetry is linear;
                # Thus, it possess mirror symmetry for any axis perpendicular
                # To the rotational axis. pymatgen does not add this symmetry
                # for all linear molecules - for example, hydrogen
                if axis == [1, 0, 0]:
                    symm_m.append(SymmOp.from_xyz_string("x,-y,z"))
                    symm_m.append(SymmOp.from_xyz_string("x,y,-z"))
                    r = SymmOp.from_xyz_string("-x,y,-z")
                elif axis == [0, 1, 0]:
                    symm_m.append(SymmOp.from_xyz_string("-x,y,z"))
                    symm_m.append(SymmOp.from_xyz_string("x,y,-z"))
                    r = SymmOp.from_xyz_string("-x,-y,z")
                elif axis == [0, 0, 1]:
                    symm_m.append(SymmOp.from_xyz_string("-x,y,z"))
                    symm_m.append(SymmOp.from_xyz_string("x,-y,z"))
                    r = SymmOp.from_xyz_string("x,-y,-z")
                # Generate a full list of SymmOps for the molecule's pointgroup
                symm_m = generate_full_symmops(symm_m, 1e-3)
                break
        # Reorient the SymmOps into mol's original frame
        if not already_oriented:
            new = []
            for op in symm_m:
                new.append(P.inverse * op * P)
            return new
        elif already_oriented:
            return symm_m
    # Handle nonlinear molecules
    else:
        pg = pga.get_pointgroup()
        symm_m = []
        for op in pg:
            symm_m.append(op)
        return symm_m

def orientation_in_wyckoff_position(
    mol,
    wyckoff_position,
    randomize=True,
    exact_orientation=False,
    already_oriented=False,
    allow_inversion=True,
    rtol = 1e-2,
):
    """
    Tests if a molecule meets the symmetry requirements of a Wyckoff position,
    and returns the valid orientations.

    Args:
        mol: a Molecule object. Orientation is arbitrary
        wyckoff_position: a pyxtal.symmetry.Wyckoff_position object
        randomize: whether or not to apply a random rotation consistent with
            the symmetry requirements
        exact_orientation: whether to only check compatibility for the provided
            orientation of the molecule. Used within general case for checking.
            If True, this function only returns True or False
        already_oriented: whether or not to reorient the principle axes
            when calling get_symmetry. Setting to True can remove redundancy,
            but is not necessary
        allow_inversion: whether or not to allow chiral molecules to be
            inverted. Should only be True if the chemical and biological
            properties of the mirror image are known to be suitable for the
            desired application

    Returns:
        a list of operations.Orientation objects which can be applied to the
        molecule while allowing it to satisfy the symmetry requirements of the
        Wyckoff position. If no orientations are found, returns False.
    """
    # For single atoms, there are no constraints
    if len(mol) == 1:
        return [Orientation([[1, 0, 0], [0, 1, 0], [0, 0, 1]], degrees=2)]

    wyckoffs = wyckoff_position.ops
    w_symm = wyckoff_position.symmetry_m
    index = wyckoff_position.index

    # Obtain the Wyckoff symmetry
    symm_w = w_symm[0]
    pga = PointGroupAnalyzer(mol)

    # Check exact orientation
    if exact_orientation is True:
        mo = deepcopy(mol)
        valid = True
        for op in symm_w:
            if not pga.is_valid_op(op):
                valid = False
        if valid is True:
            return True
        elif valid is False:
            return False

    # Obtain molecular symmetry, exact_orientation==False
    symm_m = get_symmetry(mol, already_oriented=already_oriented)
    # Store OperationAnalyzer objects for each molecular SymmOp
    chiral = True
    opa_m = []
    for op_m in symm_m:
        opa = OperationAnalyzer(op_m)
        opa_m.append(opa)
        if opa.type == "rotoinversion":
            chiral = False
        elif opa.type == "inversion":
            chiral = False

    # If molecule is chiral and allow_inversion is False,
    # check if WP breaks symmetry
    if chiral is True:
        if allow_inversion is False:
            for op in wyckoffs:
                if np.linalg.det(op.rotation_matrix) < 0:
                    printx(
                        "Warning: cannot place chiral molecule in spagegroup", priority=2,
                    )
                    return False

    # Store OperationAnalyzer objects for each Wyckoff symmetry SymmOp
    opa_w = []
    for op_w in symm_w:
        opa_w.append(OperationAnalyzer(op_w))

    # Check for constraints from the Wyckoff symmetry...
    # If we find ANY two constraints (SymmOps with unique axes), the molecule's
    # point group MUST contain SymmOps which can be aligned to these particular
    # constraints. However, there may be multiple compatible orientations of the
    # molecule consistent with these constraints
    constraint1 = None
    constraint2 = None
    for i, op_w in enumerate(symm_w):
        if opa_w[i].axis is not None:
            constraint1 = opa_w[i]
            for j, op_w in enumerate(symm_w):
                if opa_w[j].axis is not None:
                    dot = np.dot(opa_w[i].axis, opa_w[j].axis)
                    if (not np.isclose(dot, 1, rtol=rtol)) and (
                        not np.isclose(dot, -1, rtol=rtol)
                    ):
                        constraint2 = opa_w[j]
                        break
            break
    # Indirectly store the angle between the constraint axes
    if constraint1 is not None and constraint2 is not None:
        dot_w = np.dot(constraint1.axis, constraint2.axis)
    # Generate 1st consistent molecular constraints
    constraints_m = []
    if constraint1 is not None:
        for i, opa1 in enumerate(opa_m):
            if opa1.is_conjugate(constraint1):
                constraints_m.append([opa1, []])
                # Generate 2nd constraint in opposite direction
                extra = deepcopy(opa1)
                extra.axis = [opa1.axis[0] * -1, opa1.axis[1] * -1, opa1.axis[2] * -1]
                constraints_m.append([extra, []])

    # Remove redundancy for the first constraints
    list_i = list(range(len(constraints_m)))
    list_j = list(range(len(constraints_m)))
    copy = deepcopy(constraints_m)
    for i, c1 in enumerate(copy):
        if i in list_i:
            for j, c2 in enumerate(copy):
                if i > j and j in list_j and j in list_i:
                    # Check if axes are colinear
                    if np.isclose(np.dot(c1[0].axis, c2[0].axis), 1, rtol=rtol):
                        list_i.remove(j)
                        list_j.remove(j)
                    # Check if axes are symmetrically equivalent
                    else:
                        cond1 = False
                        cond2 = False
                        for opa in opa_m:
                            if opa.type == "rotation":
                                op = opa.op
                                if np.isclose(
                                    np.dot(op.operate(c1[0].axis), c2[0].axis),
                                    1,
                                    rtol=5*rtol,
                                ):
                                    cond1 = True
                                    break
                        if cond1 is True:  # or cond2 is True:
                            list_i.remove(j)
                            list_j.remove(j)
    c_m = deepcopy(constraints_m)
    constraints_m = []
    for i in list_i:
        constraints_m.append(c_m[i])

    # Generate 2nd consistent molecular constraints
    valid = list(range(len(constraints_m)))
    if constraint2 is not None:
        for i, c in enumerate(constraints_m):
            opa1 = c[0]
            for j, opa2 in enumerate(opa_m):
                if opa2.is_conjugate(constraint2):
                    dot_m = np.dot(opa1.axis, opa2.axis)
                    # Ensure that the angles are equal
                    if abs(dot_m - dot_w) < 0.02 or abs(dot_m + dot_w) < 0.02:
                        constraints_m[i][1].append(opa2)
                        # Generate 2nd constraint in opposite direction
                        extra = deepcopy(opa2)
                        extra.axis = [
                            opa2.axis[0] * -1,
                            opa2.axis[1] * -1,
                            opa2.axis[2] * -1,
                        ]
                        constraints_m[i][1].append(extra)
            # If no consistent constraints are found, remove first constraint
            if constraints_m[i][1] == []:
                valid.remove(i)
    copy = deepcopy(constraints_m)
    constraints_m = []
    for i in valid:
        constraints_m.append(copy[i])

    # Generate orientations consistent with the possible constraints
    orientations = []
    # Loop over molecular constraint sets
    for c1 in constraints_m:
        v1 = c1[0].axis
        v2 = constraint1.axis
        T = rotate_vector(v1, v2)
        # If there is only one constraint
        if c1[1] == []:
            o = Orientation(T, degrees=1, axis=constraint1.axis)
            orientations.append(o)
        else:
            # Loop over second molecular constraints
            for opa in c1[1]:
                phi = angle(constraint1.axis, constraint2.axis)
                phi2 = angle(constraint1.axis, np.dot(T, opa.axis))
                if np.isclose(phi, phi2, rtol=rtol):
                    r = np.sin(phi)
                    c = np.linalg.norm(np.dot(T, opa.axis) - constraint2.axis)
                    theta = np.arccos(1 - (c ** 2) / (2 * (r ** 2)))
                    # R = aa2matrix(constraint1.axis, theta)
                    R = Rotation.from_rotvec(theta * constraint1.axis).as_matrix()
                    T2 = np.dot(R, T)
                    a = angle(np.dot(T2, opa.axis), constraint2.axis)
                    if not np.isclose(a, 0, rtol=rtol):
                        T2 = np.dot(np.linalg.inv(R), T)
                    o = Orientation(T2, degrees=0)
                    orientations.append(o)

    # Ensure the identity orientation is checked if no constraints are found
    if constraints_m == []:
        o = Orientation(np.identity(3), degrees=2)
        orientations.append(o)

    # Remove redundancy from orientations
    list_i = list(range(len(orientations)))
    list_j = list(range(len(orientations)))
    for i, o1 in enumerate(orientations):
        if i in list_i:
            for j, o2 in enumerate(orientations):
                if i > j and j in list_j and j in list_i:
                    # m1 = o1.get_matrix(angle=0)
                    # m2 = o2.get_matrix(angle=0)
                    m1 = o1.matrix
                    m2 = o2.matrix
                    new_op = SymmOp.from_rotation_and_translation(
                        np.dot(m2, np.linalg.inv(m1)), [0, 0, 0]
                    )
                    P = SymmOp.from_rotation_and_translation(np.linalg.inv(m1), [0, 0, 0])
                    old_op = P * new_op * P.inverse
                    if pga.is_valid_op(old_op):
                        list_i.remove(j)
                        list_j.remove(j)
    copy = deepcopy(orientations)
    orientations = []
    for i in list_i:
        orientations.append(copy[i])

    #Check each of the found orientations for consistency with the Wyckoff pos.
    #If consistent, put into an array of valid orientations
    allowed = []
    for o in orientations:
        if randomize is True:
            op = o.get_op("random")
        elif randomize is False:
            op = o.get_op() #do not change
        mo = deepcopy(mol)
        mo.apply_operation(op)
        if orientation_in_wyckoff_position(mo, wyckoff_position, exact_orientation=True, randomize=False, allow_inversion=allow_inversion) is True:
            allowed.append(o)
    #Return the array of allowed orientations. If there are none, return False
    if allowed == []:
        return False
    else:
        return allowed

def make_graph(mol, tol=0.2):
    """
    make graph object for the input molecule
    """
    G = nx.Graph()
    names = {}
    for i, site in enumerate(mol._sites):
        names[i] = site.specie.value

    for i in range(len(mol)-1):
        site1 = mol.sites[i]
        for j in range(i+1, len(mol)):
            site2 = mol.sites[j]
            #remove short X-H distances
            if names[i] == "H" and names[j]=="H":
                factor = -0.5
            elif [names[i], names[j]] in [["S","S"], ["S","O"], ["O","S"], ["F","O"], ["O","F"]]:
                factor = 0.05
            elif "H" in [names[i], names[j]]:
                factor = 0.5
            else:
                factor = 1.0
            try:
                if CovalentBond.is_bonded(site1, site2, factor*tol):
                    G.add_edge(i,j)
                    #print(names[i], names[j], mol.get_distance(i, j))
            except ValueError:
                pass
    nx.set_node_attributes(G, names, 'name')

    return G

def compare_mol_connectivity(mol1, mol2, ignore_name=False):
    """
    Compare two molecules by connectivity
    """

    G1 = make_graph(mol1)
    G2 = make_graph(mol2)
    if ignore_name:
        GM = nx.isomorphism.GraphMatcher(G1, G2)
    else:
        fun = lambda n1, n2: n1['name'] == n2['name']
        GM = nx.isomorphism.GraphMatcher(G1, G2, node_match=fun)

    return GM.is_isomorphic(), GM.mapping

max_stars_repo_name: pedroantoniosantosf/PyXtal

max_issues_repo_name: pedroantoniosantosf/PyXtal

max_forks_repo_name: pedroantoniosantosf/PyXtal

reason: import numpy,from scipy,import networkx


text:
import chess
import chess.svg
import random
import datetime
import numpy as np
import timer_setup as ts
import time
import chess_bot as M1_BOT
import chess_eval_bot as BOT1
import chess_eval_bot_v2 as BOT2
import chess_eval_bot_v3 as BOT3
import chess_eval_bot_v4 as BOT4
import chess_eval_bot_v5 as BOT5

timer=ts.Timer()
timer.start()

board=chess.Board()

#print("Test prints:")
#print(board.legal_moves)
#print(list(board.legal_moves)[0])
#print()
#print()
print()

def interpret_result(result1,board):
	if result1=="1/2-1/2":
		if (board.is_seventyfive_moves()==1):
			return 2
		else:
			if board.is_insufficient_material()==1:
				return 3
			elif board.is_stalemate()==1:
				return 4
			else:
				return 0
	elif result1=="0-1":
		return -1
	elif result1=="1-0":
		return 1

def countX(lst, x):
    count = 0
    for ele in lst:
        if (ele == x):
            count = count + 1
    return count

nb_games=1
max_moves=1000
Rgames_outcomes=np.zeros((nb_games))

print_outcome=0

for game in range(0,nb_games):
	print("--- Starting game "+str(game)+" ---")
	board.reset()

	#black_BOT=BOT.random_bot()
	#black_BOT=BOT.uni_move_lvl5_bot()
	#black_BOT=BOT.human_move()
	#black_BOT=BOT.uni_move_lvl4_bot()
	white_BOT=BOT5.uni_move_eval_bot(1,board,3)
	#white_BOT=BOT.uni_move_lvl4_bot()
	black_BOT=BOT5.uni_move_eval_bot(0,board,3)

	move_count=0

	game_timer_start = time.perf_counter()
	white_time=0
	black_time=0

	while board.is_game_over()==0:
	#	print(move_count,board.move_stack)
		print(move_count,move_count%2)
		print(board)
		chess.svg.board(board, size=350)
		board
		if move_count%2==0:
			white_timer_start = time.perf_counter()
			chosen_move=white_BOT.chose_move(board)
			white_timer_stop = time.perf_counter()
			ti=white_timer_stop - white_timer_start
			print("turn took "+str(ti)+"s")
			white_time += ti
		elif move_count%2==1:
			black_timer_start = time.perf_counter()
			chosen_move=black_BOT.chose_move(board)
			black_timer_stop = time.perf_counter()
			ti=black_timer_stop - black_timer_start
			print("turn took "+str(ti)+"s")
			black_time += ti
		move_count+=1

		#possible_moves=list(board.legal_moves)
		#chosen_move=random.choice(possible_moves)
		if board.has_insufficient_material(board.turn)==1:
			Rgames_outcomes[game]=-100+200*(1-board.turn)
			break

		board.push(chosen_move)

		#if board.fullmove_number>75:
		#	print("main ", board.fullmove_number,board.outcome(),board.is_seventyfive_moves())

		if (board.outcome()!=None or board.fullmove_number>max_moves):
			#print(board.outcome(winner))
			if print_outcome==1:
				print(interpret_result(board.result()))
				print(board)

			if board.fullmove_number>max_moves:
				Rgames_outcomes[game]=2
				break
			else:
				Rgames_outcomes[game]=interpret_result(board.result(),board)

	game_timer_end = time.perf_counter()
	print(board)
	print(board.outcome(),board.result(),"game time:",game_timer_end- game_timer_start,"white time",white_time,"black time",black_time)

	output_file="save_stacks.txt"
	now = datetime.datetime.now()
	with open(output_file,"a+") as file:
		file.write("\n\n"+str(now)+"\n"+(str(board.move_stack).replace("Move.from_uci(","").replace(")","")+"\n"+"game time:"+str(game_timer_end- game_timer_start)+" white time "+str(white_time)+" black time "+str(black_time)+"\n\n"))

print(Rgames_outcomes)
nb_Black_wins=countX(Rgames_outcomes,-1)
nb_White_wins=countX(Rgames_outcomes,1)
nb_Black_wins+=countX(Rgames_outcomes,-100)
nb_White_wins+=countX(Rgames_outcomes,100)
nb_black_checkmates=countX(Rgames_outcomes,-1)
nb_white_checkmates=countX(Rgames_outcomes,1)
nb_Draws=countX(Rgames_outcomes,0)
nb_Draws+=countX(Rgames_outcomes,2)
nb_Draws+=countX(Rgames_outcomes,3)
nb_Draws+=countX(Rgames_outcomes,4)
nb_75moves=countX(Rgames_outcomes,2)
nb_insuficient_mat=countX(Rgames_outcomes,3)
nb_stalemates=countX(Rgames_outcomes,4)

print("white wins,  Draw,  Black wins")
print(nb_White_wins,nb_Draws,nb_Black_wins)
print(float(nb_White_wins)/float(nb_games),float(nb_Draws)/float(nb_games),float(nb_Black_wins)/float(nb_games))
print()
print("white CM,  Black CM")
print(nb_white_checkmates,nb_black_checkmates)
print()
print("75 moves, insuf. mat., stalemates")
print(nb_75moves, nb_insuficient_mat, nb_stalemates)

timer.stop()

ext: py

reason: import numpy


text:
import torch
import torch.nn as nn
import numpy as np

from modules.visual_extractor import VisualExtractor
from modules.encoder_decoder import EncoderDecoder

class FFAIRModel(nn.Module):
    def __init__(self, args, tokenizer):
        super(FFAIRModel, self).__init__()
        self.args = args
        self.tokenizer = tokenizer
        self.visual_extractor = VisualExtractor(args)
        self.encoder_decoder = EncoderDecoder(args, tokenizer)
        self.forward = self.forward_ffair

    def __str__(self):
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        params = sum([np.prod(p.size()) for p in model_parameters])
        return super().__str__() + '\nTrainable parameters: {}'.format(params)

    def forward_ffair(self, images, targets=None, mode='train'):
        att_feats = 0
        fc_feats = 0
        for ind in range(images.shape[1]):
            att_feats_new, fc_feats_new = self.visual_extractor(images[:, ind])
            att_feats += att_feats_new
            fc_feats += fc_feats_new
        att_feats /= images.shape[1]
        fc_feats /= images.shape[1]
        if mode == 'train':
            output = self.encoder_decoder(fc_feats, att_feats, targets, mode='forward')
        elif mode == 'sample':
            output, _ = self.encoder_decoder(fc_feats, att_feats, mode='sample')
        else:
            raise ValueError
        return output


text:
import numpy as np
import pandas as pd
from scipy.stats.stats import pearsonr, spearmanr, kendalltau
from common.exports import save_to_excel_multi

def compute_correlation(x, y, corr_type="spearman", supported_types=("pearson", "spearman", "kendall")):
    """
    Compute correlation between <x> and <y>

    This function computes the correlation between <x> and <y>. The correlation function
    is chosen using the mapping of the <supported_types> and <corr_type>. The function
    supports three types of correlation: ("pearson", "spearman", "kendall").

    Parameters
    ----------

    x : numpy ndarray
        numpy ndarray with the data

    y : numpy ndarray
        numpy ndarray with the data

    corr_type : str, optional, default "spearman"
        str with the correlation type

    supported_types : iterable, optional, default ("pearson", "spearman", "kendall")
        iterable with the supported correlation types

    Returns
    -------

    Correlation results: r, p
    """

    # Prepare the mapping
    mapping = {"pearson": pearsonr, "spearman": spearmanr, "kendall": kendalltau}

    # Compute the correlation
    if corr_type not in supported_types:
        raise ValueError(f"Corr. type {corr_type} is unsupported. Supported ones: {supported_types}")
    return mapping[corr_type.lower()](x, y)

def assess_correlation(df_feat, df_clin, computation_settings, save_as="tmp.xlsx"):
    """
    Assess the correlation between features and clinical rating scales according to the settings.

    This function computes the correlation between all features from <df_feat> and selected
    clinical rating scales from <df_clin> defined in <computation_settings>. The function
    supports three types of correlation: ("pearson", "spearman", "kendall"). The results
    can be stored to excel sheet, see: <save_as>. The following structure of the settings
    is anticipated:

        [
            {
                "scale": "class",
                "field_name": "class",
                "correlation": ("pearson", "spearman", "kendall")
            }
        ]

    Parameters
    ----------

    df_feat : pandas DataFrame with the features
        pandas DataFrame (rows=observations, cols=features)

    df_clin : pandas DataFrame with the clinical data
        pandas DataFrame (rows=observations, cols=features)

    computation_settings : list with the computation settings
        list of dicts with the computation settings

    save_as : str, optional, default "tmp.xlsx"
        str with the full-path to store the tables with the correlations
    """

    # Prepare the tables and sheets (each table and sheet pair stores the correlation for one scale)
    tables = []
    sheets = []

    # Compute the correlation between the features and clinical rating scales defined in the settings
    for setting in computation_settings:

        # Prepare the results table
        results = []

        for feat in df_feat:

            # Select the feature and clinical rating scale to analyze
            clin_data = df_clin[setting.get("field_name")]
            feat_data = df_feat[feat]

            # Select only observations with no missing values in both features and clinical scale
            clin_nans = clin_data.isnull()
            feat_nans = feat_data.isnull()
            with_nans = clin_nans | feat_nans

            clin_data = clin_data.values[np.logical_not(with_nans.values)]
            feat_data = feat_data.values[np.logical_not(with_nans.values)]

            # Prepare the dictionary of correlations
            correlations = {"feature": feat}

            # Compute the correlations and update the results
            for corr in setting.get("correlation"):
                r, p = compute_correlation(feat_data, clin_data, corr_type=corr)
                correlations.update({f"r ({corr})": round(float(r), 4)})
                correlations.update({f"p ({corr})": round(float(p), 4)})

            # Append to correlation the results
            results.append(correlations)

        # Update the resulting correlation tables and sheets
        tables.append(pd.DataFrame(results))
        sheets.append(setting.get("scale"))

    # Save the tables to the excel sheets
    save_to_excel_multi(tables, output_path=save_as, sheet_names=sheets, index=False)

    # Return the correlations
    return tables

max_stars_repo_name: BDALab/workbench

max_issues_repo_name: BDALab/workbench

max_forks_repo_name: BDALab/workbench

reason: import numpy,from scipy


text:
# Lint as: python3
"""A sensor for robot-mounted 1D lidar (laser scan)."""

from typing import Any, Iterable, Sequence, Type, Union

import gin
import gym
import numpy as np

from pybullet_envs.minitaur.envs_v2.sensors import sensor
from pybullet_envs.minitaur.vision import point_cloud_utils
from pybullet_envs.minitaur.vision import sim_camera

_MODE_TO_NUM_CHANNELS_DICT = {
    sim_camera.CameraMode.DEPTH: 1,
    sim_camera.CameraMode.RGB: 3,
    sim_camera.CameraMode.RGBD: 4,
    sim_camera.CameraMode.POINTCLOUD_WORLD_FRAME: 3,
    sim_camera.CameraMode.POINTCLOUD_ROBOT_FRAME: 3,
}

@gin.configurable
class CameraSensor(sensor.Sensor):
  """A robot-mounted sensor that returns RGBD images.

  Attributes:
    resolution: A 2-tuple (width, height) that represents the resolution of the
      camera.
    fov_degree: A floating point value that represents the field of view of the
      camera in the vertical direction. The unit is degree.
  """

  def __init__(self,
               camera_translation_from_base,
               camera_rotation_from_base,
               parent_link_id=-1,
               camera_mode=sim_camera.CameraMode.DEPTH,
               camera_update_frequency_hz=10,
               camera_stabilized=False,
               fov_degree=60,
               resolution=(32, 32),
               lower_bound: Union[float, Iterable[float]] = 0.0,
               upper_bound: Union[float, Iterable[float]] = 255.0,
               sensor_latency: Union[float, Sequence[float]] = 0.0,
               dtype: Type[Any] = np.float64,
               name="vision"):
    """Initializes the CameraSensor.

    Args:
      camera_translation_from_base: A 3-vector translation from the center of
        the specified link.
      camera_rotation_from_base: A 4-vector quaternion represents the rotation
        of the camera relative to the specified link.
      parent_link_id: The pybullet link id, where the camera is mounted on.
      camera_mode: An enum that specifies the mode that the camera operates. See
        sim_camera.CameraMode for more details.
      camera_update_frequency_hz: The frequency at which the camera will capture
        a frame.
      camera_stabilized: Whether the camera is stabilized. See
        sim_camera.MountedCamera for more details.
      fov_degree: The vertical field of view of the camera (in degree).
      resolution: A 2-tuple that represents the width and the height of the
        camera image.
      lower_bound: The lower bound of values of the camera output. It could be a
        single float or an array of floats with shape (height, width, channels).
      upper_bound: The upper bound of values of the camera output. It could be a
        single float or an array of floats with shape (height, width,
        channels)..
      sensor_latency: See base class.
      dtype: See base class.
      name: The name of the sensor.
    """
    super().__init__(
        name=name,
        sensor_latency=sensor_latency,
        interpolator_fn=sensor.closest_obs_blender)
    self._parent_link_id = parent_link_id
    self._camera_mode = camera_mode
    self._camera_translation_from_base = camera_translation_from_base
    self._camera_rotation_from_base = camera_rotation_from_base
    self.camera_update_frequency_hz = camera_update_frequency_hz
    self._time_interval_every_camera_update = (1.0 /
                                               self.camera_update_frequency_hz)
    self._camera_stabilized = camera_stabilized
    self._fov_degree = fov_degree
    self._resolution = resolution
    num_channels = _MODE_TO_NUM_CHANNELS_DICT[self._camera_mode]
    self._camera = None
    self._camera_image = None
    self._dtype = dtype
    if isinstance(lower_bound, list):
      lower_bound = np.array(lower_bound, dtype=self._dtype)
    else:
      lower_bound = lower_bound * np.ones(
          shape=(resolution[1], resolution[0], num_channels), dtype=self._dtype)
    if isinstance(upper_bound, list):
      upper_bound = np.array(upper_bound, dtype=self._dtype)
    else:
      upper_bound = upper_bound * np.ones(
          shape=(resolution[1], resolution[0], num_channels), dtype=self._dtype)
    self._observation_space = gym.spaces.Box(
        low=lower_bound, high=upper_bound, dtype=self._dtype)
    self._last_camera_image_timestamp = None

  def change_mounting_point(
      self,
      camera_translation_from_link: Sequence[float] = (0, 0, 0),
      camera_rotation_from_link: Sequence[float] = (0, 0, 0, 1),
      parent_link_id: int = -1):
    """Changes mounting point. Must be called before calls to set_robot().

    Args:
      camera_translation_from_link: A 3-vector translation from the center of
        the specified link.
      camera_rotation_from_link: A 4-vector quaternion represents the rotation
        of the camera relative to the specified link.
      parent_link_id: The pybullet link id, where the camera is mounted on.
    """
    self._parent_link_id = parent_link_id
    self._camera_translation_from_base = camera_translation_from_link
    self._camera_rotation_from_base = camera_rotation_from_link

  def set_robot(self, robot):
    super().set_robot(robot)

    self._camera = sim_camera.MountedCamera(
        pybullet_client=robot.pybullet_client,
        body_id=robot.robot_id,
        parent_link_id=self._parent_link_id,
        relative_translation=self._camera_translation_from_base,
        relative_rotation=self._camera_rotation_from_base,
        stabilized=self._camera_stabilized,
        camera_mode=self._camera_mode,
        fov_degree=self._fov_degree,
        resolution=self._resolution)

  def on_reset(self, env):
    self._env = env
    self._last_camera_image_timestamp = None
    super().on_reset(env)

  def _get_original_observation(self):
    if self._last_camera_image_timestamp is None or (
        self._robot.timestamp >= self._last_camera_image_timestamp +
        self._time_interval_every_camera_update):
      self._camera_image = self._camera.render_image().astype(self._dtype)
      self._last_camera_image_timestamp = self._robot.timestamp
    return self._robot.timestamp, self._camera_image

  def project_depth_map_to_point_cloud(self, depth_map, use_world_frame=True):
    """Convert the depth map into a 3D point cloud.

    Args:
      depth_map: A 2D numpy array with shape (height, width) which represents
        the depth map.
      use_world_frame: Whether converts the depth map into a point cloud in the
        world frame. If False, the point cloud is in the robot's local frame. If
        True, the point cloud is in the world frame if the robot's base
        position/orientation can be measured (e.g. in sim, using SLAM or mocap).

    Returns:
      A point cloud represented by a numpy array of shape (height, width, 3).
    """
    point_cloud = point_cloud_utils.distance_map_to_point_cloud(
        np.squeeze(depth_map), self.fov_degree / 180.0 * np.pi,
        depth_map.shape[1], depth_map.shape[0])
    if use_world_frame:
      point_cloud = (
          self._camera.transform_point_cloud_from_camera_to_world_frame(
              point_cloud))
    else:
      point_cloud = (
          self._camera.transform_point_cloud_from_camera_to_robot_frame(
              point_cloud))
    return point_cloud

  @property
  def resolution(self):
    return self._camera.resolution

  @property
  def fov_degree(self):
    return self._camera.fov_degree


text:
#!/usr/bin/python
"""
Simulation of CLDA control task
"""
## Imports

from db import dbfunctions
from db.tracker import dbq
from db.tracker import models

import os
import numpy as np
import multiprocessing as mp
from scipy.io import loadmat, savemat

from features.generator_features import Autostart
from features.simulation_features import SimHDF, SimTime, SimCosineTunedPointProc, SimPPFDecoderCursor, SimPPFDecoderCursorShuffled
from features.hdf_features import SaveHDF

import riglib.bmi
from riglib.bmi import train, kfdecoder, clda, ppfdecoder
from tasks import bmimultitasks, generatorfunctions as genfns
from riglib.stereo_opengl.window import WindowDispl2D
from tasks import cursor_clda_tasks

from riglib.bmi.feedback_controllers import LQRController

import pickle

# reload(kfdecoder)
# reload(ppfdecoder)
reload(clda)
# reload(riglib.bmi)
# reload(riglib.bmi.train)
from tasks import tentaclebmitasks
reload(cursor_clda_tasks)
reload(tentaclebmitasks)

import argparse
parser = argparse.ArgumentParser(description='Analyze neural control of a redundant kinematic chain')
parser.add_argument('--clean', help='', action="store_true")
parser.add_argument('--show', help='', action="store_true")
parser.add_argument('--alg', help='', action="store")
parser.add_argument('--save', help='', action="store_true")

args = parser.parse_args()

#######################
#### RML test case ####
#######################
class SimRML(SaveHDF, Autostart, SimTime, WindowDispl2D, cursor_clda_tasks.SimCLDAControlMulti):
    update_rate = 0.1
    starting_pos = (0., 0., 0.)
    rand_start = (0., 0.)
    assist_level = (0., 0.)
    target_radius = 1.8
    plant_type = 'cursor_14x14'
    win_res = (250, 140)
    def __init__(self, *args, **kwargs):
        super(SimRML, self).__init__(*args, **kwargs)
        self.batch_time = 0.1
        self.half_life  = (20.0, 20.0)
        self.starting_pos = (0., 0., 0.)
        self.assist_time = 15.

    def create_updater(self):
        self.updater = clda.KFRML(self.batch_time, self.half_life[0])

class PointProcNeuralSim(object):
    def _init_neural_encoder(self):
        from riglib.bmi import sim_neurons
        sim_encoder_fname = os.path.join(os.getenv('HOME'), 'code/bmi3d/tests/ppf', 'sample_spikes_and_kinematics_10000.mat')
        self.encoder = sim_neurons.load_ppf_encoder_2D_vel_tuning_clda_sim(sim_encoder_fname, dt=1./60)

class PointProcSmoothbatch(object):
    def create_updater(self):
        dt = 1./180
        batch_time = 1./60
        batch_size = batch_time/dt
        half_life = 120.
        rho = np.exp(np.log(0.5) / (half_life/batch_time))

        #self.updater = clda.PPFContinuousBayesianUpdater(self.decoder)
        self.updater = clda.PPFSmoothbatchSingleThread()
        self.updater.rho = rho
        self.updater.batch_time = batch_time

class PointProcContinuous(object):
    def create_updater(self):
        self.updater = clda.PPFContinuousBayesianUpdater(self.decoder)

# from tasks.cursor_clda_tasks import OFCLearner3DEndptPPF
from riglib.experiment import traits
from riglib.bmi import feedback_controllers

class SimCLDAControlMultiDispl2D_PPF(SimTime, Autostart, WindowDispl2D, SimCosineTunedPointProc, SimPPFDecoderCursorShuffled, cursor_clda_tasks.CLDAControlPPFContAdapt):
    win_res = (250, 140)
    tau = traits.Float(2.7, desc="Magic parameter for speed of OFC.")
    param_noise_scale = traits.Float(1.0, desc="Stuff")
    half_life = (0, 0)
    half_life_time = 1

    def __init__(self, *args, **kwargs):
        from riglib.bmi.state_space_models import StateSpaceEndptVel2D
        ssm = StateSpaceEndptVel2D()
        A, B, W = ssm.get_ssm_matrices(update_rate=1./180)
        Q = np.mat(np.diag([1., 1, 1, 0, 0, 0, 0]))
        R = 100*np.mat(np.diag([1., 1., 1.]))
        self.fb_ctrl = LQRController(A, B, Q, R)

        self.ssm = ssm

        super(SimCLDAControlMultiDispl2D_PPF, self).__init__(*args, **kwargs)
        self.batch_time = 1./10 #60.  # TODO 10 Hz running seems to be hardcoded somewhere
        self.assist_level = 0., 0.
        self.assist_level_time = 120.
        self.last_get_spike_counts_time = -1./60
        self.learn_flag = True

class SimCLDAControlMultiDispl2D_PPFRML(SimCLDAControlMultiDispl2D_PPF):
    half_life = (10., 120.)
    max_attempts = 1
    def create_updater(self):
        self.updater = clda.PPFRML()

if args.alg == 'RML':
    if args.save:
        te = models.TaskEntry()
        sim_subj = models.Subject.objects.using('simulation').get(name='Simulation')
        te.subject = sim_subj
        te.task = models.Task.objects.using('simulation').get(name='clda_kf_cg_rml')
        te.sequence_id = 0
        te.save(using='simulation')
    gen = cursor_clda_tasks.SimCLDAControlMulti.sim_target_seq_generator_multi(8, 100)
    task = SimRML(gen)
elif args.alg == 'PPF':
    gen = cursor_clda_tasks.SimCLDAControlMulti.sim_target_seq_generator_multi(8, 100)
    task = SimCLDAControlMultiDispl2D_PPF(gen)
elif args.alg == 'PPFRML':
    gen = cursor_clda_tasks.SimCLDAControlMulti.sim_target_seq_generator_multi(8, 100)
    task = SimCLDAControlMultiDispl2D_PPFRML(gen)
else:
    raise ValueError("Algorithm not recognized!")

self = task
task.init()
print('task init called')
task.run()

if args.save:
    task.cleanup(dbq, te.id, subject=sim_subj, dbname='simulation')

max_stars_repo_name: sgowda/brain-python-interface

max_issues_repo_name: sgowda/brain-python-interface

max_forks_repo_name: sgowda/brain-python-interface


text:
import numpy as np

from radon_server import radon_fss
from radon_server.radon_thread import RadonTransformThread

class SlowSlantStackTransform(RadonTransformThread):
    def get_matrix_ratio(self):
        return 2

    def get_reconstruct_multiply(self):
        return 1

    def get_algorithm_name(self):
        return "sss"

    def run_transform(self, image, n, variant=None):
        self.slow_slant_stack(image, n)

    def slow_slant_stack(self, image, n):
        self.radon = np.zeros((n * 2, n * 2), dtype='float64')
        k = np.fft.fftshift(np.arange(-n, n))

        kn = (np.arange(0, n) - n // 2)

        # vertical lines
        rotatedImage = np.transpose(image)
        paddedImage = np.pad(rotatedImage, n // 2, mode='constant')[n // 2:n // 2 + n, :]
        fft_matrix = np.fft.fft(paddedImage)

        for s in range(n):
            p = (float(n-s) - n / 2) / float(n / 2)
            shiftby_arr = kn * p
            shifted_fft_matrix = fft_matrix * np.exp(-1j * 2 * np.pi * k * shiftby_arr.reshape(n, 1) / (n * 2))
            skewed = np.real(np.fft.ifft(shifted_fft_matrix)).astype('float64')
            self.radon[:, s] = sum(skewed)
            self.update_progress(s, n*2)

        paddedImage = np.pad(image, n // 2, mode='constant')[n // 2:n // 2 + n, :]
        fft_matrix = np.fft.fft(paddedImage)

        # horizontal lines
        for s in range(n):
            p = (float(s) - n / 2) / float(n / 2)
            shiftby_arr = kn * p
            shifted_fft_matrix = fft_matrix * np.exp(-1j * 2 * np.pi * k * shiftby_arr.reshape(n, 1) / (n * 2))
            skewed = np.fliplr(np.real(np.fft.ifft(shifted_fft_matrix)).astype('float64'))
            self.radon[:, s + n] = np.roll(sum(skewed), 1)
            self.update_progress(n+s, n*2)

        self.radon = self.radon / (np.sqrt(2) * n)  # Normalization
        self.radon = self.radon * 2 * n

    def need_matrix(self):
        return False

    def get_matrix(self, variant, n):
        fss = radon_fss.FastSlantStackTransform(action="reconstruct", args=self.args, method=self.method)
        return fss.get_matrix(variant, n)


text:
# -*- coding: utf-8 -*-

"""
The :mod:`coclust.coclustering.base_non_diagonal_coclust` module provides a
class with common methods for non diagonal co-clustering algorithms.
"""

# Author: Francois Role <francois.role@gmail.com>
#         Stanislas Morbieu <stanislas.morbieu@gmail.com>

# License: BSD 3 clause

import numpy as np
from sklearn.base import BaseEstimator

class BaseNonDiagonalCoclust(BaseEstimator):
    def get_row_indices(self, i):
        """Give the row indices of the i’th co-cluster.

        Parameters
        ----------
        i : integer
            Index of the i'th row cluster

        Returns
        -------
        list
            list of row indices
        """
        row_indices = [index for index, label in enumerate(self.row_labels_)
                       if label == i]
        return row_indices

    def get_col_indices(self, i):
        """Give the column indices of the i’th co-cluster.

        Parameters
        ----------
        i : integer
            Index of the i'th column cluster

        Returns
        -------
        list
            list of column indices
        """
        col_indices = [index for index, label in enumerate(self.column_labels_)
                       if label == i]
        return col_indices

    def get_shape(self, i, j):
        """Give the shape of block corresponding to the i’th row cluster and
           the j'th column cluster.

        Parameters
        ----------
        i : integer
            Index of the row cluster
        j : integer
            Index of the column cluster

        Returns
        -------
        (int, int)
            (number of rows, number of columns)
        """
        row_indices = self.get_row_indices(i)
        column_indices = self.get_col_indices(i)
        return (len(row_indices), len(column_indices))

    def get_submatrix(self, m, i, j):
        """Give the submatrix corresponding to row cluster i and column cluster j.

        Parameters
        ----------
        m : X : numpy array or scipy sparse matrix
            Matrix from which the block has to be extracted
        i : integer
           index of the row cluster
        j : integer
           index of the col cluster

        Returns
        -------
        numpy array or scipy sparse matrix
            Submatrix corresponding to row cluster i and column cluster j
        """
        row_ind = np.array(self.get_row_indices(i))
        col_ind = np.array(self.get_col_indices(j))
        return m[row_ind[:, np.newaxis], col_ind]


text:
from __future__ import print_function
import argparse
import math
from math import log10
import os
import os.path
from collections import defaultdict
import itertools

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
import numpy as np

no_summary = False
try:
    from torchsummary import summary
except ModuleNotFoundError:
    no_summary = True
    print("No summary writer found")

from console_progressbar import ProgressBar

#from data import get_training_set, get_test_set
import dataset
import models
import losses
from utils import ScreenSpaceShading, initialImage

# Training settings
parser = argparse.ArgumentParser(description='Superresolution for Direct Volume Rendering')

parser.add_argument('--dataset', type=str,
                    help="Path to the HDF5 file with the dataset", required=True)
parser.add_argument('--upscale_factor', type=int, default=4, help="super resolution upscale factor")
parser.add_argument('--numberOfImages', type=int, default=-1, help="Number of images taken from the inpt dataset. Default: -1 = unlimited")

parser.add_argument('--restore', type=int, default=-1, help="Restore training from the specified run index")
parser.add_argument('--restoreEpoch', type=int, default=-1, help="In combination with '--restore', specify the epoch from which to recover. Default: last epoch")
parser.add_argument('--pretrained', type=str, default=None, help="Path to a pretrained generator")
parser.add_argument('--pretrainedDiscr', type=str, default=None, help="Path to a pretrained discriminator")

#Model parameters
parser.add_argument('--model', type=str, required=True, help="""
The superresolution model.
Supported nets: 'SubpixelNet', 'EnhanceNet', 'TecoGAN', 'RCAN'
""")
parser.add_argument('--upsample', type=str, default='bilinear', help='Upsampling for EnhanceNet: nearest, bilinear, bicubic, or pixelShuffle')
parser.add_argument('--reconType', type=str, default='residual', help='Block type for EnhanceNet: residual or direct')
parser.add_argument('--useBN', action='store_true', help='Enable batch normalization in the generator and discriminator')
parser.add_argument('--useSN', action='store_true', help='Enable spectral normalization in the generator and discriminator')
parser.add_argument('--numResidualLayers', type=int, default=10, help='Number of residual layers in the generator')
parser.add_argument('--disableTemporal', action='store_true', help='Disables temporal consistency')
parser.add_argument('--initialImage', type=str, default='zero', help="""
Specifies what should be used as the previous high res frame for the first frame of the sequence,
when no previous image is available from the previous predition.
Available options:
 - zero: fill everything with zeros
 - unshaded: Special defaults for unshaded mode: mask=-1, normal=[0,0,1], depth=0.5, ao=1
 - input: Upscale the current input image and use that as the previous high res image
   Remaining channels are filled with defaults
Default: 'input'
""")

#Loss parameters
parser.add_argument('--losses', type=str, required=True, help="""
Comma-separated list of loss functions: mse,perceptual,texture,adv.
Optinally, the weighting factor can be specified with a colon.
Example: "--losses perceptual:0.1,texture:1e2,adv:10"
""")
parser.add_argument('--perceptualLossLayers',
                    type=str,
                     # defaults found with VGGAnalysis.py
                    default='conv_1:0.026423,conv_2:0.009285,conv_3:0.006710,conv_4:0.004898,conv_5:0.003910,conv_6:0.003956,conv_7:0.003813,conv_8:0.002968,conv_9:0.002997,conv_10:0.003631,conv_11:0.004147,conv_12:0.005765,conv_13:0.007442,conv_14:0.009666,conv_15:0.012586,conv_16:0.013377',
                    help="""
Comma-separated list of layer names for the perceptual loss.
Note that the convolution layers are numbered sequentially: conv_1, conv2_, ... conv_19.
Optionally, the weighting factor can be specified with a colon: "conv_4:1.0", if omitted, 1 is used.
""")
parser.add_argument('--textureLossLayers', type=str, default='conv_1,conv_3,conv_5', help="""
Comma-separated list of layer names for the perceptual loss.
Note that the convolution layers are numbered sequentially: conv_1, conv2_, ... conv_19.
Optinally, the weighting factor can be specified with a colon: "conv_4:1.0", if omitted, 1 is used.
""")
parser.add_argument('--discriminator', type=str, default='enhanceNetLarge', help="""
Network architecture for the discriminator.
Possible values: enhanceNetSmall, enhanceNetLarge, tecoGAN
""")
#parser.add_argument('--advDiscrThreshold', type=float, default=None, help="""
#Adverserial training:
#If the cross entropy loss of the discriminator falls below that threshold, the training for the discriminator is stopped.
#Set this to zero to disable the check and use a fixed number of iterations, see --advDiscrMaxSteps, instead.
#""")
parser.add_argument('--advDiscrMaxSteps', type=int, default=2, help="""
Adverserial training:
Maximal number of iterations for the discriminator training.
Set this to -1 to disable the check.
""")
parser.add_argument('--advDiscrInitialSteps', type=int, default=None, help="""
Adverserial training:
Number of iterations for the disciriminator training in the first epoch.
Used in combination with a pretrained generator to let the discriminator catch up.
""")
parser.add_argument('--advDiscrWeightClip', type=float, default=0.01, help="""
For the Wasserstein GAN, this parameter specifies the value of the hyperparameter 'c',
the range in which the discirminator parameters are clipped.
""")
#parser.add_argument('--advGenThreshold', type=float, default=None, help="""
#Adverserial training:
#If the cross entropy loss of the generator falls below that threshold, the training for the generator is stopped.
#Set this to zero to disable the check and use a fixed number of iterations, see --advGenMaxSteps, instead.
#""")
parser.add_argument('--advGenMaxSteps', type=int, default=2, help="""
Adverserial training:
Maximal number of iterations for the generator training.
Set this to -1 to disable the check.
""")
parser.add_argument('--lossBorderPadding', type=int, default=16, help="""
Because flow + warping can't be accurately estimated at the borders of the image,
the border of the input images to the loss (ground truth, low res input, prediction)
are overwritten with zeros. The size of the border is specified by this parameter.
Pass zero to disable this padding. Default=16 as in the TecoGAN paper.
""")

parser.add_argument('--samples', type=int, required=True, help='Number of samples for the train and test dataset')
parser.add_argument('--testFraction', type=float, default=0.2, help='Fraction of test data')
parser.add_argument('--batchSize', type=int, default=16, help='training batch size')
parser.add_argument('--testBatchSize', type=int, default=16, help='testing batch size')
parser.add_argument('--testNumFullImages', type=int, default=4, help='number of full size images to test for visualization')
parser.add_argument('--nEpochs', type=int, default=2, help='number of epochs to train for')
parser.add_argument('--lr', type=float, default=0.0001, help='Learning Rate. Default=0.0001')
parser.add_argument('--lrGamma', type=float, default=0.5, help='The learning rate decays every lrStep-epochs by this factor')
parser.add_argument('--lrStep', type=int, default=500, help='The learning rate decays every lrStep-epochs (this parameter) by lrGamma factor')
parser.add_argument('--weightDecay', type=float, default=0, help="Weight decay (L2 penalty), if supported by the optimizer. Default=0")
parser.add_argument('--optim', type=str, default="Adam", help="""
Optimizers. Possible values: RMSprop, Rprop, Adam (default).
""")
parser.add_argument('--noTestImages', action='store_true', help="Don't save full size test images")

parser.add_argument('--cuda', action='store_true', help='use cuda?')
parser.add_argument('--seed', type=int, default=124, help='random seed to use. Default=124')
parser.add_argument('--logdir', type=str, default='C:/Users/Mustafa/Desktop/dvrRendering/log', help='directory for tensorboard logs')
parser.add_argument('--modeldir', type=str, default='C:/Users/Mustafa/Desktop/dvrRendering/model', help='Output directory for the checkpoints')

opt = parser.parse_args()
opt_dict = vars(opt)

if opt.cuda and not torch.cuda.is_available():
    raise Exception("No GPU found, please run without --cuda")

torch.manual_seed(opt.seed)
np.random.seed(opt.seed)

device = torch.device("cuda" if opt.cuda else "cpu")
torch.set_num_threads(4)

#########################
# RESERVE OUTPUT DIRECTORY
#########################

# Run directory
def findNextRunNumber(folder):
    files = os.listdir(folder)
    files = sorted([f for f in files if f.startswith('run')])
    if len(files)==0:
        return 0
    return int(files[-1][3:])
nextRunNumber = max(findNextRunNumber(opt.logdir), findNextRunNumber(opt.modeldir)) + 1
if opt.restore == -1:
    print('Current run: %05d'%nextRunNumber)
    runName = 'run%05d'%nextRunNumber
    logdir = os.path.join(opt.logdir, runName)
    modeldir = os.path.join(opt.modeldir, runName)
    runName = 'run%05d'%nextRunNumber
    os.makedirs(logdir)
    os.makedirs(modeldir)

#########################
# DATASETS + CHANNELS
#########################

print('===> Loading datasets')
dataset_data = dataset.dvr_dense_load_samples_hdf5(opt_dict['upscale_factor'], opt_dict, opt_dict['dataset'])
print('Dataset input images have %d channels'%dataset_data.input_channels)

input_channels = dataset_data.input_channels
assert input_channels == 4 # RGB, MASK
output_channels = dataset_data.output_channels
assert output_channels == 4 # RGB, MASK
input_channels_with_previous = input_channels + output_channels * (opt.upscale_factor ** 2)

train_set = dataset.DvrDenseDatasetFromSamples(dataset_data, False, opt.testFraction, device)
test_set = dataset.DvrDenseDatasetFromSamples(dataset_data, True, opt.testFraction, device)
test_full_set = dataset.DvrDenseDatasetFromFullImages(dataset_data, min(opt.testNumFullImages, len(dataset_data.images_low)))
training_data_loader = DataLoader(dataset=train_set, batch_size=opt.batchSize, shuffle=True)
testing_data_loader = DataLoader(dataset=test_set, batch_size=opt.testBatchSize, shuffle=False)
testing_full_data_loader = DataLoader(dataset=test_full_set, batch_size=1, shuffle=False)

#############################
# MODEL
#############################
print('===> Building model')
model = models.createNetwork(
    opt.model,
    opt.upscale_factor,
    input_channels_with_previous,
    [0,1,2,3],
    output_channels,
    opt)
model.to(device)
print('Model:')
print(model)
if not no_summary:
    summary(model,
            input_size=train_set.get_low_res_shape(input_channels_with_previous),
            device=device.type)

#############################
# LOSSES
#############################

print('===> Building losses')
criterion = losses.LossNet(
    device,
    input_channels,
    output_channels,
    train_set.get_high_res_shape()[1], #high resolution size
    opt.lossBorderPadding,
    opt)
criterion.to(device)
print('Losses:', criterion)
res = train_set.get_high_res_shape()[1]
if no_summary:
    criterion.print_summary(
        (output_channels, res, res),
        (output_channels, res, res),
        (input_channels, res, res),
        (output_channels+1, res, res),
        opt.batchSize, device)

#############################
# OPTIMIZER
#############################

print('===> Create Optimizer ('+opt.optim+')')
def createOptimizer(name, parameters, lr, weight_decay):
    if name=='Adam':
        return optim.Adam(parameters, lr=lr, weight_decay=weight_decay)
    elif name=='RMSprop':
        return optim.RMSprop(parameters, lr=lr, weight_decay=weight_decay)
    elif name=='Rprop':
        return optim.Rprop(parameters, lr=lr)
    elif name=='LBFGS':
        return optim.LBFGS(parameters, lr=lr)
    else:
        raise ValueError("Unknown optimizer "+name)
if not criterion.has_discriminator:
    adversarial_training = False
    optimizer = createOptimizer(opt.optim, model.parameters(),
                                lr=opt.lr, weight_decay=opt.weightDecay)
    #scheduler = optim.lr_scheduler.ExponentialLR(optimizer, opt.lrDecay)
    scheduler = optim.lr_scheduler.StepLR(optimizer, opt.lrStep, opt.lrGamma)
else:
    adversarial_training = True
    gen_optimizer = createOptimizer(opt.optim, model.parameters(),
                                    lr=opt.lr, weight_decay=opt.weightDecay)
    #filter(lambda p: p.requires_grad, criterion.get_discr_parameters())
    discr_optimizer = createOptimizer(
        opt.optim,
        filter(lambda p: p.requires_grad, criterion.get_discr_parameters()),
        lr=opt.lr*0.5, weight_decay=opt.weightDecay)
    #gen_scheduler = optim.lr_scheduler.ExponentialLR(gen_optimizer, opt.lrDecay)
    #discr_scheduler = optim.lr_scheduler.ExponentialLR(discr_optimizer, opt.lrDecay)
    gen_scheduler = optim.lr_scheduler.StepLR(gen_optimizer, opt.lrStep, opt.lrGamma)
    discr_scheduler = optim.lr_scheduler.StepLR(discr_optimizer, opt.lrStep, opt.lrGamma)

#############################
# PRETRAINED
#############################
if opt.pretrained is not None:
    checkpoint = torch.load(opt.pretrained)
    model.load_state_dict(checkpoint['model'].state_dict())
    #only load the state dict, not the whole model
    #this asserts that the model structure is the same
    print('Using pretrained model for the generator')
if opt.pretrainedDiscr is not None:
    assert criterion.discriminator is not None
    checkpoint = torch.load(opt.pretrainedDiscr)
    criterion.discriminator.load_state_dict(checkpoint['discriminator'])
    print('Using pretrained model for the discriminator')

#############################
# Additional Stuff: Spectral Normalization
# (placed after pretrained, because models without spectral normalization
#  can't be imported as models with normalization
#############################
if opt.useSN:
    from utils.apply_sn import apply_sn
    apply_sn(model)
    if criterion.discriminator is not None:
        apply_sn(criterion.discriminator)
    print("Spectral Normalization applied")

#############################
# OUTPUT DIRECTORIES or RESTORE
#############################

#Check for restoring
startEpoch = 1
if opt.restore != -1:
    nextRunNumber = opt.restore
    runName = 'run%05d'%nextRunNumber
    modeldir = os.path.join(opt.modeldir, runName)
    if opt.restoreEpoch == -1:
        restoreEpoch = 0
        while True:
            modelInName = os.path.join(modeldir, "model_epoch_{}.pth".format(restoreEpoch+1))
            if not os.path.exists(modelInName):
                break;
            restoreEpoch += 1
    else:
        restoreEpoch = opt.restoreEpoch

    print("Restore training from run", opt.restore,"and epoch",restoreEpoch)
    modelInName = os.path.join(modeldir, "model_epoch_{}.pth".format(restoreEpoch))
    checkpoint = torch.load(modelInName)
    #model.load_state_dict(checkpoint['state_dict'])
    model = checkpoint['model'] #Restore full model
    if adversarial_training:
        criterion.discriminator.load_state_dict(checkpoint['discriminator'])
        discr_optimizer = checkpoint['discr_optimizer']
        gen_optimizer = checkpoint['gen_optimizer']
        discr_scheduler = checkpoint['discr_scheduler']
        gen_scheduler = checkpoint['gen_scheduler']
    else:
        optimizer = checkpoint['optimizer']
        scheduler = checkpoint['scheduler']
    startEpoch = restoreEpoch

#paths
print('Current run: %05d'%nextRunNumber)
runName = 'run%05d'%nextRunNumber
logdir = os.path.join(opt.logdir, runName)
modeldir = os.path.join(opt.modeldir, runName)

optStr = str(opt);
print(optStr)
with open(os.path.join(modeldir, 'info.txt'), "w") as text_file:
    text_file.write(optStr)

#tensorboard logger
writer = SummaryWriter(logdir)
writer.add_text('info', optStr, 0)

#############################
# MAIN PART
#############################

#@profile
def trainNormal(epoch):
    epoch_loss = 0
    num_minibatch = len(training_data_loader)
    pg = ProgressBar(num_minibatch, 'Training', length=50)
    model.train()
    for iteration, batch in enumerate(training_data_loader, 0):
        pg.print_progress_bar(iteration)
        input, target = batch[0].to(device), batch[1].to(device)
        B, _, Cout, Hhigh, Whigh = target.shape
        _, _, Cin, H, W = input.shape
        assert(Cout == output_channels)
        assert(Cin == input_channels)
        assert(H == dataset_data.crop_size)
        assert(W == dataset_data.crop_size)
        assert(Hhigh == dataset_data.crop_size * opt.upscale_factor)
        assert(Whigh == dataset_data.crop_size * opt.upscale_factor)

        optimizer.zero_grad()

        previous_output = None
        loss = 0
        for j in range(dataset_data.num_frames):
            # prepare input
            if j == 0 or opt.disableTemporal:
                previous_warped = initialImage(input[:,0,:,:,:], Cout,
                                               opt.initialImage, False, opt.upscale_factor)
                # loss takes the ground truth current image as warped previous image,
                # to not introduce a bias and big loss for the first image
                previous_warped_loss = target[:,0,:,:,:]
                previous_input = F.interpolate(input[:,0,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
            else:
                previous_warped = models.VideoTools.warp_upscale(
                    previous_output,
                    flow[:, j-1, :, :, :],
                    opt.upscale_factor,
                    special_mask = True)
                previous_warped_loss = previous_warped
                previous_input = F.interpolate(input[:,j-1,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                previous_input = models.VideoTools.warp_upscale(
                    previous_input,
                    flow[:, j-1, :, :, :],
                    opt.upscale_factor,
                    special_mask = True)
            previous_warped_flattened = models.VideoTools.flatten_high(previous_warped, opt.upscale_factor)
            single_input = torch.cat((
                    input[:,j,:,:,:],
                    previous_warped_flattened),
                dim=1)
            # run generator
            prediction, _ = model(single_input)
            # evaluate cost
            input_high = F.interpolate(input[:,j,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
            loss0,_ = criterion(
                target[:,j,:,:,:],
                prediction,
                input_high,
                previous_warped_loss)
            del _
            loss += loss0
            epoch_loss += loss0.item()
            # save output
            previous_output = prediction

        loss.backward()
        optimizer.step()
    pg.print_progress_bar(num_minibatch)
    epoch_loss /= num_minibatch * dataset_data.num_frames
    print("===> Epoch {} Complete: Avg. Loss: {:.4f}".format(epoch, epoch_loss))
    writer.add_scalar('train/total_loss', epoch_loss, epoch)
    writer.add_scalar('train/lr', scheduler.get_lr()[0], epoch)
    scheduler.step()

def trainAdv_v2(epoch):
    """
    Second version of adverserial training,
    for each batch, train both discriminator and generator.
    Not full epoch for each seperately
    """
    print("===> Epoch %d Training"%epoch)
    discr_scheduler.step()
    writer.add_scalar('train/lr_discr', discr_scheduler.get_lr()[0], epoch)
    gen_scheduler.step()
    writer.add_scalar('train/lr_gen', gen_scheduler.get_lr()[0], epoch)

    disc_steps = opt.advDiscrInitialSteps if opt.advDiscrInitialSteps is not None and epoch==1 else opt.advDiscrMaxSteps
    gen_steps = opt.advGenMaxSteps

    num_minibatch = len(training_data_loader)
    model.train()
    criterion.discr_train()

    total_discr_loss = 0
    total_gen_loss = 0
    total_gt_score = 0
    total_pred_score = 0

    pg = ProgressBar(num_minibatch, 'Train', length=50)
    for iteration, batch in enumerate(training_data_loader):
        pg.print_progress_bar(iteration)
        input, flow, target = batch[0].to(device), batch[1].to(device), batch[2].to(device)
        B, _, Cout, Hhigh, Whigh = target.shape
        _, _, Cin, H, W = input.shape

        # DISCRIMINATOR
        for _ in range(disc_steps):
            discr_optimizer.zero_grad()
            gen_optimizer.zero_grad()
            loss = 0
            #iterate over all timesteps
            for j in range(dataset_data.num_frames):
                # prepare input for the generator
                if j == 0 or opt.disableTemporal:
                    previous_warped = initialImage(input[:,0,:,:,:], Cout,
                                               opt.initialImage, False, opt.upscale_factor)
                    # loss takes the ground truth current image as warped previous image,
                    # to not introduce a bias and big loss for the first image
                    previous_warped_loss = target[:,0,:,:,:]
                    previous_input = F.interpolate(input[:,0,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                else:
                    previous_warped = models.VideoTools.warp_upscale(
                        previous_output,
                        flow[:, j-1, :, :, :],
                        opt.upscale_factor,
                        special_mask = True)
                    previous_warped_loss = previous_warped
                    previous_input = F.interpolate(input[:,j-1,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                    previous_input = models.VideoTools.warp_upscale(
                        previous_input,
                        flow[:, j-1, :, :, :],
                        opt.upscale_factor,
                        special_mask = True)
                previous_warped_flattened = models.VideoTools.flatten_high(previous_warped, opt.upscale_factor)
                single_input = torch.cat((
                        input[:,j,:,:,:],
                        previous_warped_flattened),
                    dim=1)
                #evaluate generator
                with torch.no_grad():
                    prediction, _ = model(single_input)
                #prepare input for the discriminator
                gt_prev_warped = models.VideoTools.warp_upscale(
                    target[:,j-1,:,:,:],
                    flow[:, j-1, :, :, :],
                    opt.upscale_factor,
                    special_mask = True)
                #evaluate discriminator
                input_high = F.interpolate(input[:,j,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                disc_loss, gt_score, pred_score = criterion.train_discriminator(
                    input_high,
                    target[:,j,:,:,:],
                    previous_input,
                    gt_prev_warped,
                    prediction,
                    previous_warped_loss)
                loss += disc_loss
                total_gt_score += float(gt_score)
                total_pred_score += float(pred_score)
                # save output
                previous_output = torch.cat([
                    torch.clamp(prediction[:,0:1,:,:], -1, +1), # mask
                    ScreenSpaceShading.normalize(prediction[:,1:4,:,:], dim=1),
                    torch.clamp(prediction[:,4:5,:,:], 0, +1), # depth
                    torch.clamp(prediction[:,5:6,:,:], 0, +1) # ao
                    ], dim=1)
            loss.backward()
            discr_optimizer.step()
        total_discr_loss += loss.item()

        # GENERATOR
        for _ in range(disc_steps):
            discr_optimizer.zero_grad()
            gen_optimizer.zero_grad()
            loss = 0
            #iterate over all timesteps
            for j in range(dataset_data.num_frames):
                # prepare input for the generator
                if j == 0 or opt.disableTemporal:
                    previous_warped = initialImage(input[:,0,:,:,:], Cout,
                                               opt.initialImage, False, opt.upscale_factor)
                    # loss takes the ground truth current image as warped previous image,
                    # to not introduce a bias and big loss for the first image
                    previous_warped_loss = target[:,0,:,:,:]
                    previous_input = F.interpolate(input[:,0,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                else:
                    previous_warped = models.VideoTools.warp_upscale(
                        previous_output,
                        flow[:, j-1, :, :, :],
                        opt.upscale_factor,
                        special_mask = True)
                    previous_warped_loss = previous_warped
                    previous_input = F.interpolate(input[:,j-1,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                    previous_input = models.VideoTools.warp_upscale(
                        previous_input,
                        flow[:, j-1, :, :, :],
                        opt.upscale_factor,
                        special_mask = True)
                previous_warped_flattened = models.VideoTools.flatten_high(previous_warped, opt.upscale_factor)
                single_input = torch.cat((
                        input[:,j,:,:,:],
                        previous_warped_flattened),
                    dim=1)
                #evaluate generator
                prediction, _ = model(single_input)
                #evaluate loss
                input_high = F.interpolate(input[:,j,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                loss0, map = criterion(
                    target[:,j,:,:,:],
                    prediction,
                    input_high,
                    previous_input,
                    previous_warped_loss)
                loss += loss0
                # save output
                previous_output = torch.cat([
                    torch.clamp(prediction[:,0:1,:,:], -1, +1), # mask
                    ScreenSpaceShading.normalize(prediction[:,1:4,:,:], dim=1),
                    torch.clamp(prediction[:,4:5,:,:], 0, +1), # depth
                    torch.clamp(prediction[:,5:6,:,:], 0, +1) # ao
                    ], dim=1)
            loss.backward()
            gen_optimizer.step()
        total_gen_loss += loss.item()
    pg.print_progress_bar(num_minibatch)

    total_discr_loss /= num_minibatch * dataset_data.num_frames
    total_gen_loss /= num_minibatch * dataset_data.num_frames
    total_gt_score /= num_minibatch * dataset_data.num_frames
    total_pred_score /= num_minibatch * dataset_data.num_frames

    writer.add_scalar('train/discr_loss', total_discr_loss, epoch)
    writer.add_scalar('train/gen_loss', total_gen_loss, epoch)
    writer.add_scalar('train/gt_score', total_gt_score, epoch)
    writer.add_scalar('train/pred_score', total_pred_score, epoch)
    print("===> Epoch {} Complete".format(epoch))

#@profile
def test(epoch):
    avg_psnr = 0
    avg_losses = defaultdict(float)
    with torch.no_grad():
        num_minibatch = len(testing_data_loader)
        pg = ProgressBar(num_minibatch, 'Testing', length=50)
        model.eval()
        if criterion.has_discriminator:
            criterion.discr_eval()
        for iteration, batch in enumerate(testing_data_loader, 0):
            pg.print_progress_bar(iteration)
            input, target = batch[0].to(device), batch[1].to(device)
            B, _, Cout, Hhigh, Whigh = target.shape
            _, _, Cin, H, W = input.shape

            previous_output = None
            for j in range(dataset_data.num_frames):
                # prepare input
                if j == 0 or opt.disableTemporal:
                    previous_warped = initialImage(input[:,0,:,:,:], Cout,
                                               opt.initialImage, False, opt.upscale_factor)
                    # loss takes the ground truth current image as warped previous image,
                    # to not introduce a bias and big loss for the first image
                    previous_warped_loss = target[:,0,:,:,:]
                    previous_input = F.interpolate(input[:,0,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                else:
                    previous_warped = models.VideoTools.warp_upscale(
                        previous_output,
                        flow[:, j-1, :, :, :],
                        opt.upscale_factor,
                        special_mask = True)
                    previous_warped_loss = previous_warped
                    previous_input = F.interpolate(input[:,j-1,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                    previous_input = models.VideoTools.warp_upscale(
                        previous_input,
                        flow[:, j-1, :, :, :],
                        opt.upscale_factor,
                        special_mask = True)
                previous_warped_flattened = models.VideoTools.flatten_high(previous_warped, opt.upscale_factor)
                single_input = torch.cat((
                        input[:,j,:,:,:],
                        previous_warped_flattened),
                    dim=1)
                # run generator
                prediction, _ = model(single_input)
                # evaluate cost
                input_high = F.interpolate(input[:,j,:,:,:], size=(Hhigh, Whigh), mode=opt.upsample)
                loss0, loss_values = criterion(
                    target[:,j,:,:,:],
                    prediction,
                    input_high,
                    previous_warped_loss)
                avg_losses['total_loss'] += loss0.item()
                psnr = 10 * log10(1 / max(1e-10, loss_values['mse']))
                avg_losses['psnr'] += psnr
                for key, value in loss_values.items():
                    avg_losses[str(key)] += value

                # save output for next frame
                previous_output = prediction
        pg.print_progress_bar(num_minibatch)
    for key in avg_losses.keys():
        avg_losses[key] /= num_minibatch * dataset_data.num_frames
    print("===> Avg. PSNR: {:.4f} dB".format(avg_losses['psnr']))
    print("  losses:",avg_losses)
    for key, value in avg_losses.items():
        writer.add_scalar('test/%s'%key, value, epoch)

def test_images(epoch):
    def write_image(img, filename):
        out_img = img.cpu().detach().numpy()
        out_img *= 255.0
        out_img = out_img.clip(0, 255)
        out_img = np.uint8(out_img)
        writer.add_image(filename, out_img, epoch)

    with torch.no_grad():
        num_minibatch = len(testing_full_data_loader)
        pg = ProgressBar(num_minibatch, 'Test %d Images'%num_minibatch, length=50)
        model.eval()
        if criterion.has_discriminator:
            criterion.discr_eval()
        for i,batch in enumerate(testing_full_data_loader):
            pg.print_progress_bar(i)
            input = batch[0].to(device)
            B, _, Cin, H, W = input.shape
            Hhigh = H * opt.upscale_factor
            Whigh = W * opt.upscale_factor
            Cout = output_channels

            channel_mask = [0, 1, 2] #RGB

            previous_output = None
            for j in range(dataset_data.num_frames):
                # prepare input
                if j == 0 or opt.disableTemporal:
                    previous_warped = initialImage(input[:,0,:,:,:], Cout,
                                               opt.initialImage, False, opt.upscale_factor)
                else:
                    previous_warped = models.VideoTools.warp_upscale(
                        previous_output,
                        flow[:, j-1, :, :, :],
                        opt.upscale_factor,
                        special_mask = True)
                previous_warped_flattened = models.VideoTools.flatten_high(previous_warped, opt.upscale_factor)
                single_input = torch.cat((
                        input[:,j,:,:,:],
                        previous_warped_flattened),
                    dim=1)
                # run generator and cost
                prediction, residual = model(single_input)
                # write prediction image
                write_image(prediction[0, channel_mask], 'image%03d/frame%03d_prediction' % (i, j))
                # write residual image
                if residual is not None:
                    write_image(residual[0, channel_mask], 'image%03d/frame%03d_residual' % (i, j))
                # save output for next frame
                previous_output = prediction
        pg.print_progress_bar(num_minibatch)

    print("Test images sent to Tensorboard for visualization")

def checkpoint(epoch):
    model_out_path = os.path.join(modeldir, "model_epoch_{}.pth".format(epoch))
    state = {'epoch': epoch + 1, 'model': model, 'parameters':opt_dict}
    if not adversarial_training:
        state.update({'optimizer':optimizer, 'scheduler':scheduler})
    else:
        state.update({'discr_optimizer':discr_optimizer,
                      'gen_optimizer':gen_optimizer,
                      'discr_scheduler':discr_scheduler,
                      'gen_scheduler':gen_scheduler,
                      'criterion': criterion})
    torch.save(state, model_out_path)
    print("Checkpoint saved to {}".format(model_out_path))

if not os.path.exists(opt.modeldir):
    os.mkdir(opt.modeldir)
if not os.path.exists(opt.logdir):
    os.mkdir(opt.logdir)

print('===> Start Training')
if not adversarial_training:
    #test_images(0)
    for epoch in range(startEpoch, opt.nEpochs + 1):
        trainNormal(epoch)
        test(epoch)
        if (epoch < 20 or (epoch%10==0)) and not opt.noTestImages:
            test_images(epoch)
        checkpoint(epoch)
else:
    test(0)
    if not opt.noTestImages:
        test_images(0)
    for epoch in range(startEpoch, opt.nEpochs + 1):
        trainAdv_v2(epoch)
        test(epoch)
        if (epoch < 20 or (epoch%10==0)) and not opt.noTestImages:
            test_images(epoch)
        if epoch%10==0:
            checkpoint(epoch)

#writer.export_scalars_to_json(os.path.join(opt.logdir, "all_scalars.json"))
writer.close()

max_stars_repo_name: shamanDevel/AdaptiveSampling

max_issues_repo_name: shamanDevel/AdaptiveSampling

max_forks_repo_name: shamanDevel/AdaptiveSampling


text:
import torch
import torch.nn as nn
import os, shutil
import numpy as np

def repackage_hidden(h):
    """Wraps hidden states in new Tensors, to detach them from their history."""
    if isinstance(h, torch.Tensor):
        return h.detach()
    else:
        return tuple(repackage_hidden(v) for v in h)

def batchify(data, bsz, args):
    nbatch = data.size(0) // bsz
    data = data.narrow(0, 0, nbatch * bsz)
    data = data.view(bsz, -1).t().contiguous()
    print(data.size())
    data = data.cuda()
    return data

def get_batch(source, i, args, seq_len=None, evaluation=False):
    seq_len = min(seq_len if seq_len else args.bptt, len(source) - 1 - i)
    data = source[i:i+seq_len].detach() if evaluation else source[i:i+seq_len]
    target = source[i+1:i+1+seq_len]
    return data, target

def create_exp_dir(path, scripts_to_save=None):
    if not os.path.exists(path):
        os.mkdir(path)

    print('Experiment dir : {}'.format(path))
    if scripts_to_save is not None:
        os.mkdir(os.path.join(path, 'scripts'))
        for script in scripts_to_save:
            dst_file = os.path.join(path, 'scripts', os.path.basename(script))
            shutil.copyfile(script, dst_file)

def save_checkpoint(model, optimizer, epoch, path, finetune=False):
    if finetune:
        torch.save(model, os.path.join(path, 'finetune_model.pt'))
        torch.save(optimizer.state_dict(), os.path.join(path, 'finetune_optimizer.pt'))
    else:
        torch.save(model, os.path.join(path, 'model.pt'))
        torch.save(optimizer.state_dict(), os.path.join(path, 'optimizer.pt'))
    torch.save({'epoch': epoch+1}, os.path.join(path, 'misc.pt'))

def embedded_dropout(embed, words, dropout=0.1, scale=None):
    if dropout:
        mask = embed.weight.data.new().resize_((embed.weight.size(0), 1)).bernoulli_(1 - dropout).expand_as(embed.weight) / (1 - dropout)
        masked_embed_weight = mask * embed.weight
    else:
        masked_embed_weight = embed.weight
    if scale:
        masked_embed_weight = scale.expand_as(masked_embed_weight) * masked_embed_weight

    padding_idx = embed.padding_idx
    if padding_idx is None:
        padding_idx = -1
    # X = embed._backend.Embedding.apply(words, masked_embed_weight,
    #     padding_idx, embed.max_norm, embed.norm_type,
    #     embed.scale_grad_by_freq, embed.sparse
    # )
    X = nn.functional.embedding(words, masked_embed_weight,
        padding_idx, embed.max_norm, embed.norm_type,
        embed.scale_grad_by_freq, embed.sparse
    )
    return X

class LockedDropout(nn.Module):
    def __init__(self):
        super(LockedDropout, self).__init__()

    def forward(self, x, dropout=0.5):
        if not self.training or not dropout:
            return x
        m = x.data.new(1, x.size(1), x.size(2)).bernoulli_(1 - dropout)
        mask = m.div_(1 - dropout).detach()
        mask = mask.expand_as(x)
        return mask * x

def mask2d(B, D, keep_prob):
    m = torch.floor(torch.rand(B, D) + keep_prob) / keep_prob
    m = m.requires_grad_(False).cuda()
    return m