text:
##############################################################################
##
# This file is part of Sardana
##
# http://www.sardana-controls.org/
##
# Copyright 2011 CELLS / ALBA Synchrotron, Bellaterra, Spain
##
# Sardana is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
##
# Sardana is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
##
# You should have received a copy of the GNU Lesser General Public License
# along with Sardana.  If not, see <http://www.gnu.org/licenses/>.
##
##############################################################################

"""
    Macro library containning diffractometer related macros for the macros
    server Tango device server as part of the Sardana project.

"""

# TODO: use taurus instead of PyTango API e.g. read_attribute,
# write_attribute. This module is full of PyTango centric calls.

# TODO: use explicit getters to obtain Sardana elements
# (controller - getController, pseudomotor - getPseudoMotor, ...) instead of
# using getDevice. However this getter seems to accept only the elements names
# and not the full names.

__all__ = ["addreflection", "affine", "br", "ca", "caa", "ci", "computeub",
           "freeze", "getmode", "hklscan", "hscan", "kscan", "latticecal",
           "loadcrystal", "lscan", "newcrystal", "or0", "or1", "orswap",
           "pa", "savecrystal", "setaz", "setlat", "setmode", "setor0",
           "setor1", "setorn", "th2th", "ubr", "wh"]

import time
import math
import os
import re
import numpy as np

from sardana.macroserver.macro import Macro, iMacro, Type
from sardana.macroserver.macros.scan import aNscan
from sardana.macroserver.msexception import UnknownEnv

from taurus.core.util.log import Logger

logger = Logger.getLogger("MacroManager")

logger.info(
    "Diffractometer macros are at early stage. They can slightly change. Macro luppsi is not tested.")

class _diffrac:
    """Internal class used as a base class for the diffractometer macros"""

    env = ('DiffracDevice',)

    def prepare(self):

        dev_name = self.getEnv('DiffracDevice')
        self.diffrac = self.getDevice(dev_name)

        try:
            dev_name = self.getEnv('Psi')
            self.psidevice = self.getDevice(dev_name)
        except:
            pass

        motorlist = self.diffrac.motorlist

        pseudo_motor_names = []
        for motor in self.diffrac.hklpseudomotorlist:
            pseudo_motor_names.append(motor.split(' ')[0])

        self.h_device = self.getDevice(pseudo_motor_names[0])
        self.k_device = self.getDevice(pseudo_motor_names[1])
        self.l_device = self.getDevice(pseudo_motor_names[2])

        motor_list = self.diffrac.motorlist

        self.nb_motors = len(motor_list)

        try:
            self.angle_names = self.diffrac.motorroles
        except:  # Only for compatibility
            self.angle_names = []
            if self.nb_motors == 4:
                self.angle_names.append("omega")
                self.angle_names.append("chi")
                self.angle_names.append("phi")
                self.angle_names.append("tth")
            elif self.nb_motors == 6:
                self.angle_names.append("mu")
                self.angle_names.append("omega")
                self.angle_names.append("chi")
                self.angle_names.append("phi")
                self.angle_names.append("gamma")
                self.angle_names.append("delta")

        if self.nb_motors == 4:
            self.labelmotor = {'Omega': "omega",
                               'Chi': "chi", 'Phi': "phi", 'Tth': "tth"}
        elif self.nb_motors == 6:
            self.labelmotor = {'Mu': "mu", 'Theta': "omega", 'Chi': "chi",
                               'Phi': "phi", 'Gamma': "gamma",
                               'Delta': "delta"}

        prop = self.diffrac.get_property(['DiffractometerType'])
        for v in prop['DiffractometerType']:
            self.type = v

        self.angle_device_names = {}
        i = 0
        for motor in motor_list:
            self.angle_device_names[self.angle_names[i]] = motor.split(' ')[0]
            i = i + 1

    # TODO: it should not be necessary to implement on_stop methods in the
    # macros in order to stop the moveables. Macro API should provide this kind
    # of emergency stop (if the moveables are correctly reserved with the
    # getMotion method) in case of aborting a macro.
    def on_stop(self):

        for angle in self.angle_names:
            angle_dev = self.getDevice(self.angle_device_names[angle])
            angle_dev.Stop()

    def check_collinearity(self, h0, k0, l0, h1, k1, l1):

        print h0
        cpx = k0 * l1 - l0 * k1
        cpy = l0 * h1 - h0 * l1
        cpz = h0 * k1 - k0 * h1
        cp_square = math.sqrt(cpx * cpx + cpy * cpy + cpz * cpz)

        collinearity = False

        if cp_square < 0.01:
            collinearity = True

        return collinearity

    def get_hkl_ref0(self):

        reflections = []
        try:
            reflections = self.diffrac.reflectionlist
        except:
            pass

        hkl = []
        if reflections is not None:
            for i in range(1, 4):
                hkl.append(reflections[0][i])

        return hkl

    def get_hkl_ref1(self):

        reflections = []
        try:
            reflections = self.diffrac.reflectionlist
        except:
            pass

        hkl = []
        if reflections is not None:
            if len(reflections) > 1:
                for i in range(1, 4):
                    hkl.append(reflections[1][i])

        return hkl

    def fl(self, ch,
           regx=re.compile(
            '(?<![\d.])'
            '(?![1-9]\d*(?![\d.])|\d*\.\d*\.)'
            '0*(?!(?<=0)\.)'
            '([\d.]+?)'
            '\.?0*'
            '(?![\d.])'
           ),
           repl=lambda mat: mat.group(mat.lastindex)
           if mat.lastindex != 3
           else '0' + mat.group(3)):
        mat = regx.search(ch)
        if mat:
            return regx.sub(repl, ch)

class br(Macro, _diffrac):
    """Move the diffractometer to the reciprocal space coordinates given by
    H, K and L.
    If a fourth parameter is given, the combination of angles to be set is
    the correspondig to the given index. The index of the
    angles combinations are then changed."""

    param_def = [
        ['H', Type.String, None, "H value"],
        ['K', Type.String, None, "K value"],
        ['L', Type.String, None, "L value"],
        ['AnglesIndex', Type.Integer, -1, "Angles index"],
        ['FlagNotBlocking', Type.Integer,  0,
         "If 1 not block. Return without finish movement"],
        ['FlagPrinting', Type.Integer,  0,
         "If 1 printing. Used by ubr"]
    ]

    def prepare(self, H, K, L, AnglesIndex, FlagNotBlocking, FlagPrinting):
        _diffrac.prepare(self)

    def run(self, H, K, L, AnglesIndex, FlagNotBlocking, FlagPrinting):
        h_idx = 0
        k_idx = 1
        l_idx = 2

        if AnglesIndex != -1:
            sel_tr = AnglesIndex
        else:
            sel_tr = self.diffrac.selectedtrajectory

        hkl_labels = ["H", "K", "L"]

        if H in hkl_labels or K in hkl_labels or L in hkl_labels:
            try:
                q_vector = self.getEnv('Q')
            except UnknownEnv:
                self.error("Environment Q not defined. Run wh to define it")
                return
            try:
                if H in hkl_labels:
                    H = float(q_vector[h_idx])
                if K in hkl_labels:
                    K = float(q_vector[k_idx])
                if L in hkl_labels:
                    L = float(q_vector[l_idx])
            except:
                self.error("Wrong format of Q vector")
                return
        hkl_values = [float(H), float(K), float(L)]

        self.diffrac.write_attribute("computetrajectoriessim", hkl_values)

        angles_list = self.diffrac.trajectorylist[sel_tr]

        if FlagNotBlocking == 0:
            cmd = "mv"
            for name, angle in zip(self.angle_names, angles_list):
                cmd = cmd + " " + str(self.angle_device_names[name])
                cmd = cmd + " " + str(angle)
            if FlagPrinting == 1:
                cmd = "u" + cmd
            self.execMacro(cmd)
        else:
            for name, angle in zip(self.angle_names, angles_list):
                angle_dev = self.getObj(self.angle_device_names[name])
                angle_dev.write_attribute("Position", angle)

        self.setEnv('Q', [hkl_values[h_idx], hkl_values[k_idx],
                          hkl_values[l_idx], self.diffrac.WaveLength])

class ubr(Macro, _diffrac):
    """Move the diffractometer to the reciprocal space coordinates given by
    H, K and L und update.
    """

    param_def = [
        ["hh", Type.String, "Not set", "H position"],
        ["kk", Type.String, "Not set", "K position"],
        ["ll", Type.String, "Not set", "L position"],
        ['AnglesIndex', Type.Integer, -1, "Angles index"]
    ]

    def prepare(self, hh, kk, ll, AnglesIndex):
        _diffrac.prepare(self)

    def run(self, hh, kk, ll, AnglesIndex):
        if ll != "Not set":
            self.execMacro("br", hh, kk, ll, AnglesIndex, 0, 1)
        else:
            self.output("usage:  ubr H K L [Trajectory]")

class _ca(Macro, _diffrac):
    """Calculate motor positions for given H K L according to the current
    operation mode, for all trajectories or for the first one"""

    param_def = [
        ['H', Type.Float, None, "H value for the azimutal vector"],
        ['K', Type.Float, None, "K value for the azimutal vector"],
        ['L', Type.Float, None, "L value for the azimutal vector"],
        ['Trajectory', Type.Float, -1, "If -1, all trajectories"],
    ]

    def prepare(self, H, K, L, Trajectory):
        _diffrac.prepare(self)

    def run(self, H, K, L, Trajectory):

        hkl_values = [H, K, L]

        self.diffrac.write_attribute("computetrajectoriessim", hkl_values)

        if Trajectory == -1:
            start_range = 0
            end_range = len(self.diffrac.trajectorylist)
        else:
            start_range = Trajectory
            end_range = Trajectory + 1

        for i in range(int(start_range), int(end_range)):
            angles_list = self.diffrac.trajectorylist[i]
            self.output("")
            self.output("Trajectory %2d " % i)

            self.output("H K L =  %9.5f %9.5f %9.5f " %
                        (self.h_device.position, self.k_device.position,
                         self.l_device.position))

            try:
                self.output("Azimuth (Psi) = %7.5f" %
                            (self.psidevice.Position))
            except:
                self.warning(
                    "Not able to read psi. Check if environment Psi is defined")

            self.output("Wavelength =  %7.5f" % (self.diffrac.WaveLength))
            self.output("")

            str_pos = {}
            j = 0
            for name in self.angle_names:
                str_pos[name] = "%7.5f" % angles_list[j]
                j = j + 1

            if self.nb_motors == 6:
                self.output("%10s %11s %12s %11s %10s %11s" %
                            ("Delta", "Theta", "Chi", "Phi",
                             "Mu", "Gamma"))
                self.output("%10s %11s %12s %11s %10s %11s" %
                            (str_pos[self.labelmotor["Delta"]],
                             str_pos[self.labelmotor["Theta"]],
                             str_pos[self.labelmotor["Chi"]],
                             str_pos[self.labelmotor["Phi"]],
                             str_pos[self.labelmotor["Mu"]],
                             str_pos[self.labelmotor["Gamma"]]))
            elif self.nb_motors == 4:
                self.output("%10s %11s %12s %11s" %
                            ("Tth", "Omega", "Chi", "Phi"))
                self.output("%10s %11s %12s %11s" %
                            (str_pos[self.labelmotor["Tth"]],
                             str_pos[self.labelmotor["Omega"]],
                             str_pos[self.labelmotor["Chi"]],
                             str_pos[self.labelmotor["Phi"]]))
            elif self.nb_motors == 7:
                self.output("%10s %11s %12s %11s %11s %11s %11s" %
                            (self.angle_names[0],
                             self.angle_names[1],
                             self.angle_names[2],
                             self.angle_names[3],
                             self.angle_names[4],
                             self.angle_names[5],
                             self.angle_names[6]))
                self.output("%10s %11s %12s %11s %11s %11s %11s" %
                            (str_pos[self.angle_names[0]],
                             str_pos[self.angle_names[1]],
                             str_pos[self.angle_names[2]],
                             str_pos[self.angle_names[3]],
                             str_pos[self.angle_names[4]],
                             str_pos[self.angle_names[5]],
                             str_pos[self.angle_names[6]]))

class ca(Macro, _diffrac):
    """Calculate motor positions for given H K L according to the current
    operation mode (trajectory 0)."""

    param_def = [
        ['H', Type.Float, None, "H value for the azimutal vector"],
        ['K', Type.Float, None, "K value for the azimutal vector"],
        ['L', Type.Float, None, "L value for the azimutal vector"],
    ]

    def prepare(self, H, K, L):
        _diffrac.prepare(self)

    def run(self, H, K, L):

        hkl_values = [H, K, L]

        self.execMacro("_ca", H, K, L, 0)

class caa(Macro, _diffrac):
    """Calculate motor positions for given H K L according to the current
    operation mode (all trajectories)"""

    param_def = [
        ['H', Type.Float, None, "H value for the azimutal vector"],
        ['K', Type.Float, None, "K value for the azimutal vector"],
        ['L', Type.Float, None, "L value for the azimutal vector"],
    ]

    def prepare(self, H, K, L):
        _diffrac.prepare(self)

    def run(self, H, K, L):

        hkl_values = [H, K, L]

        self.execMacro("_ca", H, K, L)

class ci(Macro, _diffrac):
    """ Calculate hkl for given angle values """

    param_def = [
        ['mu', Type.Float, None, "Mu value"],
        ['theta', Type.Float, None, "Theta value"],
        ['chi', Type.Float, None, "Chi value"],
        ['phi', Type.Float, None, "Phi value"],
        ['gamma', Type.Float, -999, "Gamma value"],
        ['delta', Type.Float, -999, "Delta value"],
        ['omega_t', Type.Float, -999, "Omega_t value"],
    ]

    def prepare(self, mu, theta, chi, phi, gamma, delta, omega_t):
        _diffrac.prepare(self)

    def run(self, mu, theta, chi, phi, gamma, delta, omega_t):

        if delta == -999 and self.nb_motors == 6:
            self.error("Six angle values are need as argument")
        elif omega_t == -999 and self.nb_motors == 7:
            msg = ("Seven angle values are need as argument (omega_t is "
                   "missed - last argument)")
            self.error(msg)
        else:
            if self.nb_motors != 7:
                angles = [mu, theta, chi, phi, gamma, delta]
            else:
                angles = [omega_t, mu, theta, chi, phi, gamma, delta]

            self.diffrac.write_attribute("computehkl", angles)

            hkl_values = self.diffrac.computehkl

            self.output("h %f k %f l %f" %
                        (hkl_values[0], hkl_values[1], hkl_values[2]))

class pa(Macro, _diffrac):
    """Prints information about the active diffractometer."""

    suffix = ("st", "nd", "rd", "th")

    def prepare(self):
        _diffrac.prepare(self)

    def run(self):

        str_type = "Eulerian 6C"
        if self.type == 'E4CV':
            str_type = "Eulerian 4C Vertical"
        elif self.type == 'E4CH':
            str_type = "Eulerian 4C Horizontal"
        elif self.type == 'K6C':
            str_type = "Kappa 6C"
        elif self.type == 'K4CV':
            str_type = "Kappa 4C Vertical"

        self.output("%s Geometry (%s), %s" %
                    (str_type, self.type, self.diffrac.enginemode))
        #self.output("Sector %s" % "[ToDo]")
        self.output("")

        reflections = self.diffrac.reflectionlist

        nb_ref = 0
        if reflections is not None:
            for ref in reflections:
                if nb_ref < len(self.suffix):
                    sf = self.suffix[nb_ref]
                else:
                    sf = self.suffix[3]
                self.output("  %d%s Reflection (index %d): " %
                            (nb_ref + 1, sf, ref[0]))
                #self.output("    Affinement, Relevance : %d %d" % (ref[4], ref[5]))
                if len(ref) > 10:
                    self.output("    %s %s %s %s %s %s: %s %s %s %s %s %s" %
                                (self.angle_names[5], self.angle_names[1],
                                 self.angle_names[2], self.angle_names[3],
                                 self.angle_names[4], self.angle_names[0],
                                 _diffrac.fl(self, str(ref[11])),
                                 _diffrac.fl(self, str(ref[7])),
                                 _diffrac.fl(self, str(ref[8])),
                                 _diffrac.fl(self, str(ref[9])),
                                 _diffrac.fl(self, str(ref[10])),
                                 _diffrac.fl(self, str(ref[6]))))
                else:
                    self.output("    %s %s %s %s: %s %s %s %s" %
                                (self.angle_names[0], self.angle_names[1],
                                 self.angle_names[2], self.angle_names[3],
                                 _diffrac.fl(self, str(ref[6])),
                                 _diffrac.fl(self, str(ref[7])),
                                 _diffrac.fl(self, str(ref[8])),
                                 _diffrac.fl(self, str(ref[9]))))
                nb_ref = nb_ref + 1
                self.output(" %33s  %s %s %s" % ("H K L =",
                                                 _diffrac.fl(
                                                     self, str(ref[1])),
                                                 _diffrac.fl(
                                                     self, str(ref[2])),
                                                 _diffrac.fl(
                                                     self, str(ref[3]))))
                self.output("")

#        self.output("")
        self.output("  Lattice Constants (lengths / angles):")
        self.output("%32s = %s %s %s / %s %s %s" % ("real space", self.diffrac.a,
                                                    self.diffrac.b, self.diffrac.c, _diffrac.fl(
                                                        self, str(self.diffrac.alpha)),
                                                    _diffrac.fl(self, str(
                                                        self.diffrac.beta)),
                                                    _diffrac.fl(self, str(
                                                        self.diffrac.gamma))))

        self.output("")
        self.output("  Azimuthal reference:")
        self.output("%34s %s %s %s " %
                    ("H K L =", _diffrac.fl(self, str(self.diffrac.psirefh)), _diffrac.fl(self, str(self.diffrac.psirefk)), _diffrac.fl(self, str(self.diffrac.psirefl))))

        self.output("")
        self.output("  Lambda = %s" % (self.diffrac.WaveLength))

        lst = self.diffrac.ubmatrix
        self.output("  UB-Matrix")
        self.output("  %15g %15g %15g" % (lst[0][0], lst[0][1], lst[0][2]))
        self.output("  %15g %15g %15g" % (lst[1][0], lst[1][1], lst[1][2]))
        self.output("  %15g %15g %15g" % (lst[2][0], lst[2][1], lst[2][2]))

class wh(Macro, _diffrac):
    """Show principal axes and reciprocal space positions.

    Prints the current reciprocal space coordinates (H K L) and the user
    positions of the principal motors. Depending on the diffractometer
    geometry, other parameters such as the angles of incidence and
    reflection (ALPHA and BETA) and the incident wavelength (LAMBDA)
    may be displayed."""

    def prepare(self):
        _diffrac.prepare(self)

    def run(self):

        self.output("")
        self.output("Engine: %s" % self.diffrac.engine)
        self.output("")
        self.output("Mode: %s" % self.diffrac.enginemode)

        self.output("")
        self.output("%s %s %3s %9.5f %9.5f %9.5f " %
                    ("H", "K", "L = ", self.h_device.position,
                     self.k_device.position, self.l_device.position))

        if self.diffrac.psirefh == -999:
            self.output("")
        else:
            self.output("%8s %9.5f %9.5f %9.5f " %
                        ("Ref   = ", self.diffrac.psirefh, self.diffrac.psirefk, self.diffrac.psirefl))

            psirefh_in = self.diffrac.psirefh
            psirefk_in = self.diffrac.psirefk
            psirefl_in = self.diffrac.psirefl
            engine_restore = self.diffrac.engine
            mode_restore = self.diffrac.enginemode

            self.diffrac.write_attribute("engine", "psi")

            psirefh_psi = self.diffrac.psirefh
            psirefk_psi = self.diffrac.psirefk
            psirefl_psi = self.diffrac.psirefl

            self.diffrac.write_attribute("engine", engine_restore)
            self.diffrac.write_attribute("enginemode", mode_restore)

            if psirefh_in != psirefh_psi or psirefk_in != psirefk_psi or psirefl_in != psirefl_psi:
                self.warning(
                    "Psiref vector missmatch. Calculated value corresponds to:")
                self.warning("%8s %9.5f %9.5f %9.5f " %
                             ("Ref   = ", psirefh_psi, psirefk_psi, psirefl_psi))
                self.warning("Use setaz for setting it consistently")

        try:
            self.output("%s %7.5f" % (
                "Azimuth (Psi - calculated) = ", self.psidevice.Position))
        except:
            self.warning(
                "Not able to read psi. Check if environment Psi is defined")

        parameter_names = self.diffrac.modeparametersnames

        if parameter_names is not None:
            i = 0
            for par in parameter_names:
                if par == "psi":
                    parameter_values = self.diffrac.modeparametersvalues
                    self.info("%s %7.5f" %
                              ("Azimuth (Psi - set) = ", parameter_values[i]))
                i = i + 1

        self.output("%s %7.5f" % ("Wavelength = ", self.diffrac.WaveLength))
        self.output("")

        if self.nb_motors == 6:
            str_pos1 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Delta"]]).Position
            str_pos2 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Theta"]]).Position
            str_pos3 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Chi"]]).Position
            str_pos4 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Phi"]]).Position
            str_pos5 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Mu"]]).Position
            str_pos6 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Gamma"]]).Position
            self.output("%10s %11s %12s %11s %10s %11s" %
                        ("Delta", "Theta", "Chi", "Phi", "Mu", "Gamma"))
            self.output("%10s %11s %12s %11s %10s %11s" %
                        (str_pos1, str_pos2, str_pos3, str_pos4, str_pos5,
                         str_pos6))
        elif self.nb_motors == 4:
            str_pos1 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Tth"]]).Position
            str_pos2 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Omega"]]).Position
            str_pos3 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Chi"]]).Position
            str_pos4 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.labelmotor["Phi"]]).Position
            self.output("%10s %11s %12s %11s" %
                        ("Tth", "Omega", "Chi", "Phi"))
            self.output("%10s %11s %12s %11s" %
                        (str_pos1, str_pos2, str_pos3, str_pos4))
        elif self.nb_motors == 7:
            str_pos1 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.angle_names[0]]).Position
            str_pos2 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.angle_names[1]]).Position
            str_pos3 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.angle_names[2]]).Position
            str_pos4 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.angle_names[3]]).Position
            str_pos5 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.angle_names[4]]).Position
            str_pos6 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.angle_names[5]]).Position
            str_pos7 = "%7.5f" % self.getDevice(
                self.angle_device_names[self.angle_names[6]]).Position
            self.output("%10s %11s %12s %11s %11s %11s %11s" %
                        (self.angle_names[0],
                         self.angle_names[1],
                         self.angle_names[2],
                         self.angle_names[3],
                         self.angle_names[4],
                         self.angle_names[5],
                         self.angle_names[6]))
            self.output("%10s %11s %12s %11s %11s %11s %11s" %
                        (str_pos1,
                         str_pos2,
                         str_pos3,
                         str_pos4,
                         str_pos5,
                         str_pos6,
                         str_pos7))

        self.setEnv('Q', [self.h_device.position, self.k_device.position,
                          self.l_device.position, self.diffrac.WaveLength])

class freeze(Macro, _diffrac):
    """ Set psi value for psi constant modes """

    param_def = [
        ['parameter', Type.String, None, "Parameter to freeze"],
        ['value',     Type.Float,  None, "Value to be frozen"]
    ]

    def prepare(self, parameter, value):
        _diffrac.prepare(self)

    def run(self, parameter, value):

        if parameter == "psi":
            engine_restore = self.diffrac.engine
            mode_restore = self.diffrac.enginemode

            if mode_restore != "psi_constant_vertical" and mode_restore != "psi_constant_horizontal":
                self.warning(
                    "Psi frozen to set value. But current mode is not set to psi_constant_vertical or psi_constant_horizontal ")

            self.diffrac.write_attribute("engine", "hkl")
            self.diffrac.write_attribute("enginemode", "psi_constant_vertical")
            parameter_values = self.diffrac.modeparametersvalues
            parameter_values[3] = value
            self.diffrac.write_attribute(
                "modeparametersvalues", parameter_values)
            self.diffrac.write_attribute(
                "enginemode", "psi_constant_horizontal")
            parameter_values = self.diffrac.modeparametersvalues
            parameter_values[3] = value
            self.diffrac.write_attribute(
                "modeparametersvalues", parameter_values)

            self.diffrac.write_attribute("engine", engine_restore)
            self.diffrac.write_attribute("enginemode", mode_restore)

        else:
            self.error("Only implemented for parameter psi. Nothing done")

class setmode(iMacro, _diffrac):
    """Set operation mode."""

    param_def = [
        ['new_mode', Type.Integer, -1, "Mode to be set"]
    ]

    def prepare(self, new_mode):
        _diffrac.prepare(self)

    def run(self, new_mode):

        modes = self.diffrac.enginemodelist

        if new_mode == -1:
            self.output("Available modes:")
            imode = 1
            old_mode = self.diffrac.read_attribute("enginemode").value
            for mode in modes:
                if mode == old_mode:
                    def_mode = imode
                self.output(" %d -> %s " % (imode, mode))
                imode = imode + 1
            old_mode = self.diffrac.read_attribute("enginemode").value
            self.output("")
            a = self.input("Your choice?", default_value=def_mode)
            imode = 1
            for mode in modes:
                if imode == int(a):
                    def_mode = imode
                imode = imode + 1

            self.diffrac.write_attribute("enginemode", modes[def_mode - 1])
            self.output("")
            self.output("Now using %s mode" % modes[def_mode - 1])
            return

        if new_mode > len(modes):
            self.output(
                "Wrong index mode -> only from 1 to %d allowed:" % len(modes))
            imode = 1
            for mode in modes:
                self.output(" %d -> %s " % (imode, mode))
                imode = imode + 1
            return
        else:
            self.diffrac.write_attribute("enginemode", modes[new_mode - 1])
            self.output("Now using %s mode" % modes[new_mode - 1])

            self.execMacro('savecrystal')

class getmode(Macro, _diffrac):
    """Get operation mode."""

    def prepare(self):
        _diffrac.prepare(self)

    def run(self):

        self.output(self.diffrac.enginemode)

class setlat(iMacro, _diffrac):
    """Set the crystal lattice parameters a, b, c, alpha, beta and gamma
       for the currently active diffraction pseudo motor controller."""

    param_def = [
        ['a', Type.Float, -999, "Lattice 'a' parameter"],
        ['b', Type.Float, -999, "Lattice 'b' parameter"],
        ['c', Type.Float, -999, "Lattice 'c' parameter"],
        ['alpha', Type.Float, -999, "Lattice 'alpha' parameter"],
        ['beta',  Type.Float, -999, "Lattice 'beta' parameter"],
        ['gamma', Type.Float, -999, "Lattice 'gamma' parameter"]
    ]

    def prepare(self, a, b, c, alpha, beta, gamma):
        _diffrac.prepare(self)

    def run(self, a, b, c, alpha, beta, gamma):

        if gamma == -999:
            a = self.diffrac.a
            b = self.diffrac.b
            c = self.diffrac.c
            alpha = self.diffrac.alpha
            beta = self.diffrac.beta
            gamma = self.diffrac.gamma
            self.output("")
            self.output("Enter real space lattice parameters:")
            a = self.input(" Lattice a?", default_value=a,
                           data_type=Type.String)
            b = self.input(" Lattice b?", default_value=b,
                           data_type=Type.String)
            c = self.input(" Lattice c?", default_value=c,
                           data_type=Type.String)
            alpha = self.input(" Lattice alpha?",
                               default_value=alpha, data_type=Type.String)
            beta = self.input(" Lattice beta?",
                              default_value=beta, data_type=Type.String)
            gamma = self.input(" Lattice gamma?",
                               default_value=gamma, data_type=Type.String)
            self.output("")
            self.diffrac.write_attribute("a", float(a))
            self.diffrac.write_attribute("b", float(b))
            self.diffrac.write_attribute("c", float(c))
            self.diffrac.write_attribute("alpha", float(alpha))
            self.diffrac.write_attribute("beta", float(beta))
            self.diffrac.write_attribute("gamma", float(gamma))
        else:
            self.diffrac.write_attribute("a", a)
            self.diffrac.write_attribute("b", b)
            self.diffrac.write_attribute("c", c)
            self.diffrac.write_attribute("alpha", alpha)
            self.diffrac.write_attribute("beta", beta)
            self.diffrac.write_attribute("gamma", gamma)

        self.execMacro('computeub')

class or0(Macro, _diffrac):
    """Set primary orientation reflection."""

    param_def = [
        ['H', Type.Float, None, "H value"],
        ['K', Type.Float, None, "K value"],
        ['L', Type.Float, None, "L value"],
    ]

    def prepare(self, H, K, L):
        _diffrac.prepare(self)

    def run(self, H, K, L):

        # Check collinearity

        hkl_ref1 = _diffrac.get_hkl_ref1(self)
        if len(hkl_ref1) > 1:
            check = _diffrac.check_collinearity(
                self, H, K, L, hkl_ref1[0], hkl_ref1[1], hkl_ref1[2])
            if check:
                self.warning(
                    "Can not orient: or0 %9.5f %9.5f %9.5f are parallel to or1" % (H, K, L))
                return

        values = [0, H, K, L]
        self.diffrac.write_attribute("SubstituteReflection", values)

        self.execMacro('computeub')

class or1(Macro, _diffrac):
    """Set secondary orientation reflection."""

    param_def = [
        ['H', Type.Float, None, "H value"],
        ['K', Type.Float, None, "K value"],
        ['L', Type.Float, None, "L value"],
    ]

    def prepare(self, H, K, L):
        _diffrac.prepare(self)

    def run(self, H, K, L):

        # Check collinearity

        hkl_ref0 = _diffrac.get_hkl_ref0(self)
        if len(hkl_ref0) > 1:
            check = _diffrac.check_collinearity(
                self, hkl_ref0[0], hkl_ref0[1], hkl_ref0[2], H, K, L)
            if check:
                self.warning(
                    "Can not orient: or0 is parallel to or1 %9.5f %9.5f %9.5f" % (H, K, L))
                return

        values = [1, H, K, L]
        self.diffrac.write_attribute("SubstituteReflection", values)

        self.execMacro('computeub')

class setor0(Macro, _diffrac):
    """Set primary orientation reflection choosing hkl and angle values.
       Run it without any argument to see the order real positions"""

    param_def = [
        ['H', Type.Float, -999, "H value"],
        ['K', Type.Float, -999, "K value"],
        ['L', Type.Float, -999, "L value"],
        ['ang1', Type.Float, -999, "Real position"],
        ['ang2', Type.Float, -999, "Real position"],
        ['ang3', Type.Float, -999, "Real position"],
        ['ang4', Type.Float, -999, "Real position"],
        ['ang5', Type.Float, -999, "Real position"],
        ['ang6', Type.Float, -999, "Real position"],
        ['ang7', Type.Float, -999, "Real position"],
    ]

    def prepare(self, H, K, L, ang1, ang2, ang3, ang4, ang5, ang6, ang7):
        _diffrac.prepare(self)

    def run(self, H, K, L, ang1, ang2, ang3, ang4, ang5, ang6, ang7):
        setorn, pars = self.createMacro(
            "setorn", 0, H, K, L, ang1, ang2, ang3, ang4, ang5, ang6, ang7)

        self.runMacro(setorn)

class setor1(Macro, _diffrac):
    """Set secondary orientation reflection choosing hkl and angle values.
       Run it without any argument to see the order real positions"""

    param_def = [
        ['H', Type.Float, -999, "H value"],
        ['K', Type.Float, -999, "K value"],
        ['L', Type.Float, -999, "L value"],
        ['ang1', Type.Float, -999, "Real position"],
        ['ang2', Type.Float, -999, "Real position"],
        ['ang3', Type.Float, -999, "Real position"],
        ['ang4', Type.Float, -999, "Real position"],
        ['ang5', Type.Float, -999, "Real position"],
        ['ang6', Type.Float, -999, "Real position"],
        ['ang7', Type.Float, -999, "Real position"],
    ]

    def prepare(self, H, K, L, ang1, ang2, ang3, ang4, ang5, ang6, ang7):
        self.output("setor1 prepare")
        self.output(ang3)
        self.output(ang7)
        _diffrac.prepare(self)

    def run(self, H, K, L, ang1, ang2, ang3, ang4, ang5, ang6, ang7):
        setorn, pars = self.createMacro(
            "setorn", 1, H, K, L, ang1, ang2, ang3, ang4, ang5, ang6, ang7)

        self.runMacro(setorn)

class setorn(iMacro, _diffrac):
    """Set orientation reflection indicated by the index.
       Run it without any argument to see the order of the angles to be set"""

    param_def = [
        ['ref_id', Type.Integer, None, "reflection index (starting at 0)"],
        ['H', Type.Float, -999, "H value"],
        ['K', Type.Float, -999, "K value"],
        ['L', Type.Float, -999, "L value"],
        ['ang1', Type.Float, -999, "Real position"],
        ['ang2', Type.Float, -999, "Real position"],
        ['ang3', Type.Float, -999, "Real position"],
        ['ang4', Type.Float, -999, "Real position"],
        ['ang5', Type.Float, -999, "Real position"],
        ['ang6', Type.Float, -999, "Real position"],
        ['ang7', Type.Float, -999, "Real position"],
    ]

    def prepare(self, ref_id, H, K, L, ang1, ang2, ang3, ang4, ang5, ang6,
                ang7):
        _diffrac.prepare(self)

    def run(self, ref_id, H, K, L, ang1, ang2, ang3, ang4, ang5, ang6, ang7):

        if H == -999.0:
            msg = "Order of the real motor positions to be given as argument:"
            self.output(msg)
            for el in self.angle_names:
                self.output(el)
            return

        if ((ang6 == -999 and self.nb_motors == 6)
                or (ang4 == -999 and self.nb_motors == 4)
                or (ang7 == -999 and self.nb_motors == 7)):
            reflections = []
            try:
                reflections = self.diffrac.reflectionlist
            except:
                pass
            tmp_ref = {}
            hkl_names = ["h", "k", "l"]
            if reflections is not None:
                if len(reflections) > ref_id:
                    for i in range(1, 4):
                        tmp_ref[hkl_names[i - 1]] = reflections[ref_id][i]
                    for i in range(6, 6 + self.nb_motors):
                        tmp_ref[
                            self.angle_names[i - 6]] = reflections[ref_id][i]
                else:
                    for i in range(0, 3):
                        tmp_ref[hkl_names[i]] = 0
                    for i in range(0, self.nb_motors):
                        tmp_ref[self.angle_names[i]] = 0
            else:
                for i in range(0, 3):
                    tmp_ref[hkl_names[i]] = 0
                for i in range(0, self.nb_motors):
                    tmp_ref[self.angle_names[i]] = 0

            self.output("")
            if ref_id == 0:
                ref_txt = "primary-reflection"
            elif ref_id == 1:
                ref_txt = "secondary-reflection"
            else:
                ref_txt = "reflection " + str(ref_id)

            self.output("Enter %s angles" % ref_txt)
            angles_to_set = []
            for el in self.angle_names:
                angles_to_set.append(
                    float(self.input(el+"?",
                                     default_value=tmp_ref[el],
                                     data_type=Type.String)
                          )
                )

            self.output("")
            self.output("Enter %s HKL coordinates" % ref_txt)
            H = float(self.input(" H?", default_value=tmp_ref[
                      "h"], data_type=Type.String))
            K = float(self.input(" K?", default_value=tmp_ref[
                      "k"], data_type=Type.String))
            L = float(self.input(" L?", default_value=tmp_ref[
                      "l"], data_type=Type.String))
            self.output("")
        else:
            angles = [ang1, ang2, ang3, ang4, ang5, ang6, ang7]
            angles_to_set = []
            for i in range(0, self.nb_motors):
                angles_to_set.append(angles[i])

        # Check collinearity

        if ref_id == 0:
            hkl_ref = _diffrac.get_hkl_ref1(self)
        if ref_id == 1:
            hkl_ref = _diffrac.get_hkl_ref0(self)
        if ref_id < 2:
            if len(hkl_ref) > 1:
                check = _diffrac.check_collinearity(
                    self, hkl_ref[0], hkl_ref[1], hkl_ref[2], H, K, L)
                if check:
                    self.warning(
                        "Can not orient: ref0 is parallel to ref1 %9.5f %9.5f %9.5f" % (H, K, L))
                    return

        # Set reflection

        values = [ref_id, H, K, L]
        self.diffrac.write_attribute("SubstituteReflection", values)

        values = []
        values.append(ref_id)
        for el in angles_to_set:
            values.append(el)

        self.diffrac.write_attribute("AdjustAnglesToReflection", values)

        # Recompute u

        self.execMacro('computeub')

class setaz(iMacro, _diffrac):
    """ Set hkl values of the psi reference vector"""

    param_def = [
        ['PsiH', Type.Float, -999, "H value of psi reference vector"],
        ['PsiK', Type.Float, -999, "K value of psi reference vector"],
        ['PsiL', Type.Float, -999, "L value of psi reference vector"],
    ]

    def prepare(self, PsiH, PsiK, PsiL):
        _diffrac.prepare(self)

    def run(self, PsiH, PsiK, PsiL):
        engine_restore = self.diffrac.engine
        mode_restore = self.diffrac.enginemode

        if PsiL == -999:
            self.diffrac.write_attribute("engine", "hkl")
            self.diffrac.write_attribute("enginemode", "psi_constant_vertical")
            azh = self.diffrac.read_attribute("psirefh").value
            azk = self.diffrac.read_attribute("psirefk").value
            azl = self.diffrac.read_attribute("psirefl").value
            self.output("")
            self.output("Enter azimuthal reference H K L:")
            a1 = self.input(" Azimuthal H?", default_value=azh,
                            data_type=Type.String)
            a2 = self.input(" Azimuthal K?", default_value=azk,
                            data_type=Type.String)
            a3 = self.input(" Azimuthal L?", default_value=azl,
                            data_type=Type.String)
            PsiH = float(a1)
            PsiK = float(a2)
            PsiL = float(a3)

        self.diffrac.write_attribute("engine", "hkl")
        self.diffrac.write_attribute("enginemode", "psi_constant_vertical")
        self.diffrac.write_attribute("psirefh", PsiH)
        self.diffrac.write_attribute("psirefk", PsiK)
        self.diffrac.write_attribute("psirefl", PsiL)
        self.diffrac.write_attribute("enginemode", "psi_constant_horizontal")
        self.diffrac.write_attribute("psirefh", PsiH)
        self.diffrac.write_attribute("psirefk", PsiK)
        self.diffrac.write_attribute("psirefl", PsiL)
        self.diffrac.write_attribute("engine", "psi")
        self.diffrac.write_attribute("psirefh", PsiH)
        self.diffrac.write_attribute("psirefk", PsiK)
        self.diffrac.write_attribute("psirefl", PsiL)

        self.diffrac.write_attribute("engine", engine_restore)
        self.diffrac.write_attribute("enginemode", mode_restore)
        self.execMacro('savecrystal')

class computeub(Macro, _diffrac):
    """ Compute UB matrix with reflections 0 and 1 """

    def prepare(self):
        _diffrac.prepare(self)

    def run(self):

        reflections = self.diffrac.reflectionlist
        if reflections is not None:
            if len(reflections) > 1:
                self.output("Computing UB with reflections 0 and 1")
                values = [0, 1]
                self.diffrac.write_attribute("ComputeUB", values)
                self.execMacro('savecrystal')
            else:
                self.warning("UB can not be computed. Only one reflection")
        else:
            self.warning("UB can not be computed. No reflection")

class addreflection(Macro, _diffrac):
    """ Add reflection at the botton of reflections list """

    param_def = [
        ['H', Type.Float, None, "H value"],
        ['K', Type.Float, None, "K value"],
        ['L', Type.Float, None, "L value"],
        ['affinement', Type.Float, -999., "Affinement"]
    ]

    def prepare(self, H, K, L, affinement):
        _diffrac.prepare(self)

    def run(self, H, K, L, affinement):

        values = [H, K, L]
        if affinement != -999.:
            values.append(affinement)

        self.diffrac.write_attribute("AddReflection", values)

class affine(Macro, _diffrac):
    """Affine current crystal.
    Fine tunning of lattice parameters and UB matrix based on
    current crystal reflections. Reflections with affinement
    set to 0 are not used. A new crystal with the post fix
    (affine) is created and set as current crystal"""

    def prepare(self):
        _diffrac.prepare(self)

    def run(self):

        self.diffrac.write_attribute("AffineCrystal", 0)

class orswap(Macro, _diffrac):
    """Swap values for primary and secondary vectors."""

    def prepare(self):
        _diffrac.prepare(self)

    def run(self):

        self.diffrac.write_attribute("SwapReflections01", 0)
        self.output("Orientation vectors swapped.")
        self.execMacro('computeub')

class newcrystal(iMacro, _diffrac):
    """ Create a new crystal (if it does not exist) and select it. """

    param_def = [
        ['crystal_name',  Type.String, "", 'Name of the crystal to add and select']
    ]

    def prepare(self, crystal_name):
        _diffrac.prepare(self)

    def run(self, crystal_name):

        crystal_list = self.diffrac.crystallist

        to_add = 1
        i = 1
        if crystal_name == "":
            crystal_name = self.diffrac.crystal
            self.output("Available crystals:")
            for crystal in crystal_list:
                self.output("(%s) %s" % (i, crystal))
                if crystal_name == crystal:
                    iselname = crystal
                i = i + 1
            a = self.input("New crystal?", default_value=iselname,
                           data_type=Type.String)
            try:
                a1 = int(a)
                i = 1
                for crystal in crystal_list:
                    if a1 == i:
                        a = crystal
                    i = i + 1
                if a1 > i - 1:
                    a = iselname
            except:
                pass

            if a != iselname:
                crystal_name = a
            else:
                crystal_name = iselname

        for crystal in crystal_list:
            if crystal_name == crystal:
                to_add = 0

        if to_add:
            self.diffrac.write_attribute("addcrystal", crystal_name)

        self.diffrac.write_attribute("crystal", crystal_name)

        self.output("")
        self.output("Crystal selected: %s " % crystal_name)

        if to_add:
            a = self.input(" Lattice a?", default_value=5.43,
                           data_type=Type.String)
            b = self.input(" Lattice b?", default_value=5.43,
                           data_type=Type.String)
            c = self.input(" Lattice c?", default_value=5.43,
                           data_type=Type.String)
            alpha = self.input(" Lattice alpha?",
                               default_value=90, data_type=Type.String)
            beta = self.input(" Lattice beta?",
                              default_value=90, data_type=Type.String)
            gamma = self.input(" Lattice gamma?",
                               default_value=90, data_type=Type.String)
            self.output("")
            self.diffrac.write_attribute("a", float(a))
            self.diffrac.write_attribute("b", float(b))
            self.diffrac.write_attribute("c", float(c))
            self.diffrac.write_attribute("alpha", float(alpha))
            self.diffrac.write_attribute("beta", float(beta))
            self.diffrac.write_attribute("gamma", float(gamma))

class hscan(aNscan, Macro, _diffrac):
    "Scan h axis"

    param_def = [
        ['start_pos',  Type.Float,   None, 'Scan start position'],
        ['final_pos',  Type.Float,   None, 'Scan final position'],
        ['nr_interv',  Type.Integer, None, 'Number of scan intervals'],
        ['integ_time', Type.Float,   None, 'Integration time'],
    ]

    def prepare(self, start_pos, final_pos, nr_interv, integ_time):
        _diffrac.prepare(self)
        aNscan._prepare(self, [self.h_device], [start_pos],
                        [final_pos], nr_interv, integ_time)

class kscan(aNscan, Macro, _diffrac):
    "Scan k axis"

    param_def = [
        ['start_pos',  Type.Float,   None, 'Scan start position'],
        ['final_pos',  Type.Float,   None, 'Scan final position'],
        ['nr_interv',  Type.Integer, None, 'Number of scan intervals'],
        ['integ_time', Type.Float,   None, 'Integration time'],
    ]

    def prepare(self, start_pos, final_pos, nr_interv, integ_time):
        _diffrac.prepare(self)
        aNscan._prepare(self, [self.k_device], [start_pos],
                        [final_pos], nr_interv, integ_time)

class lscan(aNscan, Macro, _diffrac):
    "Scan l axis"

    param_def = [
        ['start_pos',  Type.Float,   None, 'Scan start position'],
        ['final_pos',  Type.Float,   None, 'Scan final position'],
        ['nr_interv',  Type.Integer, None, 'Number of scan intervals'],
        ['integ_time', Type.Float,   None, 'Integration time'],
    ]

    def prepare(self, start_pos, final_pos, nr_interv, integ_time):
        _diffrac.prepare(self)
        aNscan._prepare(self, [self.l_device], [start_pos],
                        [final_pos], nr_interv, integ_time)

class hklscan(aNscan, Macro, _diffrac):
    "Scan h k l axes"

    param_def = [
        ['h_start_pos',  Type.Float,   None, 'Scan h start position'],
        ['h_final_pos',  Type.Float,   None, 'Scan h final position'],
        ['k_start_pos',  Type.Float,   None, 'Scan k start position'],
        ['k_final_pos',  Type.Float,   None, 'Scan k final position'],
        ['l_start_pos',  Type.Float,   None, 'Scan l start position'],
        ['l_final_pos',  Type.Float,   None, 'Scan l final position'],
        ['nr_interv',  Type.Integer, None, 'Number of scan intervals'],
        ['integ_time', Type.Float,   None, 'Integration time'],
    ]

    def prepare(self, h_start_pos, h_final_pos, k_start_pos, k_final_pos, l_start_pos, l_final_pos, nr_interv, integ_time):
        _diffrac.prepare(self)
        aNscan._prepare(self, [self.h_device, self.k_device, self.l_device],
                        [h_start_pos, k_start_pos, l_start_pos], [h_final_pos,
                                                                  k_final_pos,
                                                                  l_final_pos],
                        nr_interv, integ_time)

class th2th(Macro):
    """th2th - scan:

    Relative scan around current position in del and th with d_th=2*d_delta
    """

    param_def = [
        ['rel_start_pos',  Type.Float,   -999, 'Scan start position'],
        ['rel_final_pos',  Type.Float,   -999, 'Scan final position'],
        ['nr_interv',  Type.Integer, -999, 'Number of scan intervals'],
        ['integ_time', Type.Float,   -999, 'Integration time']
    ]

    def run(self, rel_start_pos, rel_final_pos, nr_interv, integ_time):

        if ((integ_time != -999)):
            motor_del = self.getObj("del")
            motor_th = self.getObj("th")
            pos_del = motor_del.getPosition()
            pos_th = motor_th.getPosition()
            scan = self.d2scan(motor_del, rel_start_pos, rel_final_pos,
                               motor_th, rel_start_pos / 2, rel_final_pos / 2,
                               nr_interv, integ_time)
        else:
            self.output(
                "Usage:   th2th tth_start_rel tth_stop_rel intervals time")

count_scan = 1

class HookPars:
    pass

def hook_pre_move(self, hook_pars):
    global count_scan

    self.execMacro('freeze', 'psi', hook_pars.psi_save + +
                   hook_pars.angle_start + (count_scan - 1) * hook_pars.angle_interv)
    self.execMacro('ubr', hook_pars.h, hook_pars.k, hook_pars.l)

    count_scan = count_scan + 1

class luppsi(Macro, _diffrac):
    """psi scan:

    Relative scan psi angle

    [TODO] Still not tested
    """

    param_def = [
        ['rel_start_angle',  Type.Float,   -999, 'Relative start scan angle'],
        ['rel_final_angle',  Type.Float,   -999, 'Relative final scan angle'],
        ['nr_interv',  Type.Integer, -999, 'Number of scan intervals'],
        ['integ_time', Type.Float,   -999, 'Integration time']
    ]

    def prepare(self, H, K, L, AnglesIndex):
        _diffrac.prepare(self)

    def run(self, rel_start_angle, rel_final_angle, nr_interv, integ_time):

        global count_scan
        count_scan = 1

        if ((integ_time != -999)):
            self.diffrac.write_attribute("engine", "hkl")
            self.diffrac.write_attribute("enginemode", "psi_constant_vertical")
            h = self.h_device.position
            k = self.k_device.position
            l = self.l_device.position

            psi_positions = []

            try:
                psi_save = self.psidevice.Position
            except:
                self.error(
                    "Not able to read psi. Check if environment Psi is defined")
                return

            angle_interv = abs(rel_final_angle - rel_start_angle) / nr_interv

            # Construct scan macro

            self.output(self.psidevice.alias())
            psi_motor = self.getMotor(self.psidevice.alias())
            self.output(psi_motor)

            macro, pars = self.createMacro('dscan %s %f %f %d %f ' %
                                           (self.psidevice.alias(), rel_start_angle, rel_final_angle, nr_interv, integ_time))

            # Parameters for scan hook function

            hook_pars = HookPars()
            hook_pars.psi_save = psi_save
            hook_pars.angle_interv = angle_interv
            hook_pars.angle_start = rel_start_angle
            hook_pars.h = h
            hook_pars.k = k
            hook_pars.l = l
            f = lambda: hook_pre_move(self, hook_pars)
            macro.hooks = [
                (f, ["pre-move"]),
            ]

            # Start the scan

            self.runMacro(macro)

            # Return to start position

            self.info("Return to start position " + str(psi_save))
            self.execMacro('freeze', 'psi', psi_save)
            self.execMacro('ubr', h, k, l)
            self.psidevice.write_attribute("Position", psi_save)

        else:
            self.output(
                "Usage:  luppsi rel_startangle  rel_stopangle n_intervals time")

class savecrystal(Macro, _diffrac):
    """
         Save crystal information to file
    """

    def prepare(self):
        _diffrac.prepare(self)

    def run(self):

        self.info("Be aware: changes in crystal file format are still possible.")
        self.diffrac.write_attribute("SaveCrystal", 1)

class loadcrystal(iMacro, _diffrac):
    """
         Load crystal information from file
    """

    def prepare(self):
        _diffrac.prepare(self)

    def run(self):
        self.info("Be aware: changes in crystal file format are still possible.")
        active_dir = ""
        try:
            files = os.listdir(os.path.expanduser('~') + '/crystals/')
            active_dir = os.path.expanduser('~') + '/crystals/'
        except:
            self.output(
                "Directory for loading files %s/crystals does not exist" % os.path.expanduser('~'))
            newdir = self.input("Type new directory")
            try:
                files = os.listdir(newdir)
                active_dir = newdir
            except:
                self.output("New directory %s not found" % newdir)
                return

        res = filter(lambda x: x.endswith('.txt'), files)
        if len(res) == 0:
            self.output("No crystals available in set directory. Nothing done")
            return

        i = 1
        for filename in res:
            filename = filename.split('.')[0]
            self.output("(%s) %s" % (i, filename))
            i = i + 1
        a0 = self.input("Your choice? ")
        try:
            a1 = int(a0)
            i = 1
            for filename in res:
                if i == int(a0) and i < len(res) + 1:
                    file = filename
                i = i + 1
            if a1 < len(res) + 1:
                self.output("")
                self.output("File to load %s" % active_dir + file)
            else:
                self.output("Input out of range!")

            self.diffrac.write_attribute("loadcrystal", active_dir + file)
            self.diffrac.read_attribute("loadcrystal")

        except:
            if a0 != "":
                self.output("Wrong input!")
            else:
                self.output("An input file has to be given. Nothing done")

class latticecal(iMacro, _diffrac):
    """
        Calibrate lattice parameters a, b or c to current 2theta value
    """

    param_def = [
        ["parameter", Type.String, "", "Parameter"],
    ]

    def prepare(self, parameter):
        _diffrac.prepare(self)

    def run(self, parameter):
        if parameter != "":
            if parameter == "a" or parameter == "b" or parameter == "c":
                if parameter == "a":
                    a0 = self.diffrac.a
                    self.output("Old lattice parameter %s = %s" %
                                (parameter, a0))
                    h0 = self.h_device.position
                    h1 = round(h0)
                    a1 = h1 / h0 * a0
                    self.output("New lattice parameter %s = %s" %
                                (parameter, a1))
                    self.diffrac.write_attribute("a", a1)
                if parameter == "b":
                    a0 = self.diffrac.b
                    self.output("Old lattice parameter %s = %s" %
                                (parameter, a0))
                    h0 = self.k_device.position
                    h1 = round(h0)
                    a1 = h1 / h0 * a0
                    self.output("New lattice parameter %s = %s" %
                                (parameter, a1))
                    self.diffrac.write_attribute("b", a1)
                if parameter == "c":
                    a0 = self.diffrac.c
                    self.output("Old lattice parameter %s = %s" %
                                (parameter, a0))
                    h0 = self.l_device.position
                    h1 = round(h0)
                    a1 = h1 / h0 * a0
                    self.output("New lattice parameter %s = %s" %
                                (parameter, a1))
                    self.diffrac.write_attribute("c", a1)

                self.execMacro('computeub')

            else:
                self.output("Lattice parameter a, b or c")

        else:
            self.output(
                "Calibration of lattice parameters a, b or c to current 2theta value")
            self.output("usage:  latticecal parameter")

class _blockprintmove(Macro, _diffrac):
    """This macro is internal and reserved to the hkl infrastucture
    """
    param_def = [
        ['flagprint', Type.Integer, 0, '1 for printing']
    ]

    def prepare(self, flagprint):
        _diffrac.prepare(self)

    def run(self, flagprint):

        moving = 1
        tmp_dev = {}
        for angle in self.angle_names:
            tmp_dev[angle] = self.getDevice(self.angle_device_names[angle])
        while(moving):
            moving = 0
            for angle in self.angle_names:
                # TODO: For Taurus 4 / Taurus 3 compatibility
                if hasattr(mot_dev, "stateObj"):
                    angle_state = tmp_dev[angle].stateObj.read().rvalue
                else:
                    angle_state = tmp_dev[angle].state()
                if angle_state == 6:
                    moving = 1
            if flagprint == 1:
                self.outputBlock(" %7.5f  %7.5f  %7.5f" % (
                    self.h_device.position, self.k_device.position, self.l_device.position))
                self.flushOutput()
            self.checkPoint()
            time.sleep(1.0)
        if flagprint == 1:
            self.outputBlock(" %7.5f  %7.5f  %7.5f" % (
                self.h_device.position, self.k_device.position, self.l_device.position))
            self.flushOutput()

class _diff_scan(Macro):
    """Perfoms an scan keeping the data for further analysis/moves.
    This macro is internal and reserved to the hkl infrastucture.
    """
    param_def = [
        ['motor',      Type.Motor,   None, 'Motor to move'],
        ['start_pos',  Type.Float,   None, 'Scan start position'],
        ['final_pos',  Type.Float,   None, 'Scan final position'],
        ['nr_interv',  Type.Integer, None, 'Number of scan intervals'],
        ['integ_time', Type.Float,   None, 'Integration time'],
        ['channel',    Type.ExpChannel,   None, 'Channel to analize']
    ]

    def run(self, motor, start_pos, final_pos, nr_interv, integ_time, channel):

        ascan, pars = self.createMacro(
            "ascan", motor, start_pos, final_pos, nr_interv, integ_time)
        self.runMacro(ascan)

        channel_fullname = channel.getFullName()
        motor_name = motor.getName()

        arr_data = []
        arr_motpos = []
        for record in ascan.data.records:
            record_data = record.data
            arr_data.append(record_data[channel_fullname])
            arr_motpos.append(record_data[motor_name])

        # Find motor position corresponding to the maximum of channel values
        idx_max = np.argmax(arr_data)
        pos_max = arr_motpos[idx_max]

        self.output("Position to move")
        self.output(pos_max)

max_stars_repo_name: schooft/sardana

max_issues_repo_name: schooft/sardana

max_forks_repo_name: schooft/sardana


text:
import argparse
import json
import multiprocessing as mp
import numpy as np
import os
import pathlib
import pickle
import sys
import time
import torch
from collections import defaultdict
from queue import Empty, Queue

from utils.io import (
    load_image_size,
    load_intrinsics,
    load_mesh,
    load_poses,
    write_images,
)
from utils.labeling import (
    AreaFilter,
    LabelFilter,
    LabelMerger,
)
from utils.linalg import (
    decompose_mat4,
    from_hom,
    make_M_from_tqs,
    to_hom,
)
from utils.raster import Rasterizer

def load_models(d: dict, cad_root: str):
    trs = d['trs']
    scan_to_s2c = make_M_from_tqs(
        trs['translation'],
        trs['rotation'],
        trs['scale']
    )
    s2c_to_scan = np.linalg.inv(scan_to_s2c)

    models = d['aligned_models']
    meshes = []
    model_to_scans = []
    cache = {}

    for model in models:
        cat_id = model['catid_cad']
        cad_id = model['id_cad']

        obj_path = os.path.abspath(
            '{}/{}/{}/models/model_normalized.obj'.
            format(cad_root, cat_id, cad_id)
        )
        if obj_path in cache:
            mesh = cache[obj_path]
            mesh = mesh.clone()  # Copy
        else:
            mesh = load_mesh(obj_path)
            cache[obj_path] = mesh.clone()

        model_trs = model['trs']
        model_to_s2c = make_M_from_tqs(
            model_trs['translation'],
            model_trs['rotation'],
            model_trs['scale']
        )

        model_to_scan = s2c_to_scan @ model_to_s2c
        model_to_scans.append(model_to_scan)

        # mesh.transform(model_to_scan)
        points = mesh.verts_list()[0]
        points = torch.from_numpy(
            from_hom(
                to_hom(points.numpy()) @ model_to_scan.T
            )
        )
        mesh = type(mesh)(verts=[points], faces=mesh.faces_list())

        meshes.append(mesh)

    return meshes, models, model_to_scans

def render_scene(
    args: argparse.Namespace,
    meshes: list,
    models: list,
    model_to_scans: list,
    scene_id: str,
    taxonomy: list,
    label_ids: list,
    matches: dict
):
    alignments = defaultdict(lambda: [])
    images = {}

    scene_dir = os.path.join(args.scan_root, scene_id)

    vertices = [mesh.verts_list()[0] for mesh in meshes]
    mesh_centers = [
        to_hom(.5 * (vert.max(0)[0] + vert.min(0)[0]).numpy())
        for vert in vertices
    ]
    del vertices

    poses = load_poses(scene_dir)
    intrinsics = load_intrinsics(scene_dir)
    height, width = load_image_size(scene_dir)

    categories = set(t['shapenet'] for t in taxonomy)

    rasterizer = Rasterizer(width, height, intrinsics)
    area_filter = AreaFilter(ratio=args.min_ratio)
    label_filter = None
    if not args.label_merge_2d:
        label_filter = LabelFilter(
            scene_dir=scene_dir,
            match_ids=matches['to_remove'][scene_id],
            num_models=len(meshes),
            taxonomy=taxonomy,
            label_ids=label_ids,
            label_names=matches['id_to_nyu'][scene_id]
        )

    for name, pose in sorted(poses.items(), key=lambda pair: int(pair[0])):
        # Reset the rendering process
        rasterizer.clear_models()

        if np.any(np.isnan(pose)) or np.any(np.isinf(pose)):
            print('skipping {}: pose is {}'.format(name, pose))
            continue

        scan_to_camera = np.linalg.inv(pose)

        alignment_key = '{}/{}'.format(scene_id, name)

        for idx, (model_to_scan, mesh, model, center) in enumerate(zip(
            model_to_scans,
            meshes,
            models,
            mesh_centers
        )):
            # Reserve zero to background
            idx = idx + 1

            '''# Filter out objects whose centers are outside the image
            center_in_image = True
            center = intrinsics @ scan_to_camera @ center
            x, y, z = from_hom(center).tolist()
            if z <= 0:
                center_in_image = False
            else:
                row, col = y / z, x / z
                in_image = ((0 <= row < height) and (0 <= col < width))
                if not in_image:
                    center_in_image = False
            '''

            # Filter out non-benchmark classes
            if not args.all_s2c and model['catid_cad'] not in categories:
                continue

            # Transform the mesh
            # mesh = type(mesh)(mesh)  # copy
            # mesh.transform(scan_to_camera)
            point_list = mesh.verts_list()
            assert len(point_list) == 1
            points = from_hom(to_hom(point_list[0].numpy()) @ scan_to_camera.T)
            mesh = type(mesh)(
                verts=[torch.from_numpy(points).float()],
                faces=mesh.faces_list()
            )

            model_to_camera = scan_to_camera @ model_to_scan
            camera_to_model = np.linalg.inv(model_to_camera)

            rasterizer.add_mesh(mesh, idx, camera_to_model)

            t, q, s = decompose_mat4(model_to_camera)

             # Apply center filter in camera space
            center_in_image = True
            x, y, z = from_hom(intrinsics @ to_hom(t)).tolist()
            if z <= 0:
                center_in_image = False
            else:
                row, col = y / z, x / z
                center_in_image = (0 <= row < height) and (0 <= col < width)

            alignments[alignment_key].append({
                't': t.tolist(),
                'q': q.tolist(),
                's': s.tolist(),
                'catid_cad': model['catid_cad'],
                'id_cad': model['id_cad'],
                'id': idx,
                'center_in_image': center_in_image,
                'sym': model['sym']
            })

        rasterizer.rasterize()

        try:
            # print(name)
            if args.all_s2c:
                instance_image = rasterizer.read_idx()
                new_instances = []
            elif args.label_merge_2d:
                instance_image, new_instances = LabelMerger(
                    scene_dir=scene_dir,
                    image_name=name,
                    taxonomy=taxonomy,
                    label_ids=label_ids,
                    idx=rasterizer.read_idx(),
                    alignments=alignments[alignment_key],
                    num_models=len(meshes)
                ).result
            else:
                instance_image, new_instances = label_filter(
                    idx=rasterizer.read_idx(),
                    image_name=name,
                    alignments=alignments[alignment_key]
                )

        except Exception as e:
            print('Skipping {}/{} due to an exception {}'.format(
                scene_id, name, str(e)
            ))
            continue

        alignments[alignment_key].extend(new_instances)
        instance_image = area_filter(instance_image)
        instance_mask = (instance_image != 0)
        noc = (
            rasterizer.read_noc() * instance_mask[:, :, None]
            if args.render_nocs
            else None
        )

        write_images(
            output_scene_dir=os.path.join(args.output_image_dir, scene_id),
            output_image_name=name,
            noc=noc,
            depth=rasterizer.read_depth() * instance_mask,
            idx=instance_image,
            noc_offset=args.noc_offset,
            noc_scale=args.noc_scale,
            depth_scale=args.depth_scale
        )

        images[alignment_key] = {
            'height': height,
            'width': width,
            'intrinsics': intrinsics[:3, :3].tolist()
        }

    return dict(alignments), images

def worker(
    args: argparse.Namespace,
    data_queue: mp.Queue,
    taxonomy,
    label_ids,
    matches,
    process_idx,
    lock=None
):
    pid = os.getpid()
    alignments = {}
    images = {}

    while True:
        try:
            # datum = data_queue.get_nowait()
            datum = data_queue.get(timeout=20)
        except Empty:
            break

        if lock is not None:
            lock.acquire()
        print('Worker {} with pid {} is processing {} ({} / {})'.format(
            process_idx, pid, datum['id_scan'],
            datum['index'], datum['total_size']
        ))
        if lock is not None:
            lock.release()

        scene_id = datum['id_scan']

        meshes, models, model_to_scans = load_models(datum, args.cad_root)
        scene_alignments, scene_images = render_scene(
            args,
            meshes,
            models,
            model_to_scans,
            scene_id,
            taxonomy,
            label_ids,
            matches
        )
        alignments.update(scene_alignments)
        images.update(scene_images)

    output_data = {
        'alignments': alignments,
        'images': images
    }

    # output_data is too big for a queue, this solution is not optimal but safe
    if process_idx > 0:
        with open('{}.s2c.local'.format(process_idx), 'wb') as f:
            pickle.dump(output_data, f)
    else:
        return output_data

def main(args):
    parser = argparse.ArgumentParser()
    parser.add_argument('--s2c_root', type=str, required=True)
    parser.add_argument('--cad_root', type=str, required=True)
    parser.add_argument('--scan_root', type=str, required=True)
    parser.add_argument('--output_json_dir', type=str, required=True)
    parser.add_argument('--output_image_dir', type=str, required=True)
    parser.add_argument('--match_json', type=str, default='')

    parser.add_argument('--noc_offset', type=float, default=1.)
    parser.add_argument('--noc_scale', type=int, default=10000)
    parser.add_argument('--depth_scale', type=int, default=1000)
    parser.add_argument('--min_ratio', type=float, default=0)
    parser.add_argument('--label_merge_2d', type=int, default=1)
    parser.add_argument('--render_nocs', type=int, default=0)
    parser.add_argument('--all_s2c', type=int, default=1)
    parser.add_argument('--taxonomy_9', type=int, default=1)

    parser.add_argument('--num_workers', type=int, default=0)
    parser.add_argument('--limit', type=int, default=-1)
    parser.add_argument('--skip', type=int, default=-1)
    parser.add_argument('--start_file', type=str, default='')

    args = parser.parse_args(args)
    print(args)

    pathlib.Path(args.output_image_dir).mkdir(exist_ok=True, parents=True)
    pathlib.Path(args.output_json_dir).mkdir(exist_ok=True, parents=True)

    images = {}
    alignments = {}
    if args.start_file != '':
        with open(args.start_file) as f:
            start = json.load(f)
            images = start['images']
            alignments = start['alignments']

    assert args.match_json != '' or args.label_merge_2d, \
        'Argparse: You must provide match_json file or set label_merge_2d to 1'

    def clean_up():
        for fname in os.listdir(os.getcwd()):
            if '.s2c.local' in fname:
                os.unlink(os.path.join(os.getcwd(), fname))

    clean_up()

    with open(os.path.join(args.s2c_root, 'full_annotations.json')) as f:
        data = json.load(f)
        if args.skip > 0:
            data = data[args.skip:]
        if args.limit > 0:
            limit = min(args.limit, len(data))
            data = data[:limit]
        used_ids = set(k.split('/')[0] for k in images.keys())
        data = [d for d in data if d['id_scan'] not in used_ids]

    path = '../metadata/scan2cad_taxonomy{}.json'.format(
        '_9' if args.taxonomy_9 else ''
    )
    print('Reading {}...'.format(path))
    with open(os.path.abspath(path)) as f:
        taxonomy = json.load(f)
    del path

    with open(os.path.abspath('../metadata/labelids_all.txt')) as f:
        label_ids = []
        for line in f:
            parts = line.split()
            label_ids.append((int(parts[0]), ' '.join(parts[1:])))

    matches = None
    if not args.label_merge_2d:
        with open(args.match_json) as f:
            matches = json.load(f)

    distributed = args.num_workers > 0

    if distributed:
        data_queue = mp.Queue(maxsize=len(data))
        lock = mp.Lock()
        workers = [
            mp.Process(
                target=worker,
                args=(
                    args,
                    data_queue,
                    taxonomy,
                    label_ids,
                    matches,
                    i + 1,
                    lock,
                )
            ) for i in range(args.num_workers)
        ]
    else:
        data_queue = Queue(maxsize=len(data))

    for i, d in enumerate(data):
        d['index'] = i
        d['total_size'] = len(data)
        scene_id = d['id_scan']
        # if scene_id not in ('scene0697_01', 'scene0356_02', 'scene0678_01'):
        #    continue
        if scene_id != 'scene0470_00':
            continue
        data_queue.put(d)

    if distributed:
        for process in workers:
            process.start()

        for process in workers:
            process.join()
    else:
        local_data = worker(args, data_queue, taxonomy, label_ids, matches, 0)

    print('Collecting data...')

    def update(local_data):
        images.update(local_data['images'])
        alignments.update(local_data['alignments'])

    if distributed:
        for fname in os.listdir(os.getcwd()):
            if '.s2c.local' in fname:
                path = os.path.join(os.getcwd(), fname)
                with open(path, 'rb') as f:
                    local_data = pickle.load(f)
                del path
                update(local_data)
    else:
        update(local_data)

    output_data = {
        'alignments': alignments,
        'images': images,
        'config': {
            'all_s2c': args.all_s2c,
            'noc_scale': args.noc_scale,
            'noc_offset': args.noc_offset,
            'depth_scale': args.depth_scale,
            'min_ratio': args.min_ratio,
            'taxonomy_9': args.taxonomy_9
        }
    }
    output_path = os.path.join(
        args.output_json_dir, 'scan2cad_image_alignments.json'
    )
    print('Writing alignments to {}...'.format(output_path))
    with open(output_path, 'w') as f:
        json.dump(output_data, f)

    print('Cleaning up...')
    clean_up()

if __name__ == '__main__':
    start = time.time()
    main(sys.argv[1:])
    end = time.time()
    print('Elapsed time: {} seconds'.format(end - start))

ext: py

max_stars_repo_name: cangumeli/ROCA

max_issues_repo_name: cangumeli/ROCA

max_forks_repo_name: cangumeli/ROCA

reason: import numpy


text:
import os
import timeit
import warnings

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.metrics import mean_squared_error, mean_absolute_error

from fedot.core.data.data import InputData
from fedot.core.data.data_split import train_test_data_setup
from fedot.core.data.multi_modal import MultiModalData
from fedot.core.pipelines.node import PrimaryNode, SecondaryNode
from fedot.core.pipelines.pipeline import Pipeline
from fedot.core.repository.dataset_types import DataTypesEnum
from fedot.core.repository.tasks import Task, TsForecastingParams, TaskTypesEnum
from fedot.core.utils import fedot_project_root

warnings.filterwarnings('ignore')
np.random.seed(2020)

def make_forecast(pipeline, train: InputData, predict: InputData,
                  train_exog: InputData, predict_exog: InputData):
    """
    Function for predicting values in a time series

    :return predicted_values: numpy array, forecast of model
    """

    # Fit it
    start_time = timeit.default_timer()

    second_node_name = 'exog_ts'

    if train_exog is None:
        second_node_name = 'data_source_ts/2'
        train_exog = train
        predict_exog = predict

    train_dataset = MultiModalData({
        'data_source_ts/1': train,
        second_node_name: train_exog})

    predict_dataset = MultiModalData({
        'data_source_ts/1': predict,
        second_node_name: predict_exog})

    pipeline.fit_from_scratch(train_dataset)
    amount_of_seconds = timeit.default_timer() - start_time

    print(f'\nIt takes {amount_of_seconds:.2f} seconds to train pipeline\n')

    # Predict
    predicted_values = pipeline.predict(predict_dataset)
    predicted_values = predicted_values.predict

    return predicted_values

def run_exogenous_experiment(path_to_file, len_forecast=250, with_exog=True,
                             with_visualisation=True) -> None:
    """ Function with example how time series forecasting can be made with using
    exogenous features

    :param path_to_file: path to the csv file with dataframe
    :param len_forecast: forecast length
    :param with_exog: is it needed to make prediction with exogenous time series
    :param with_visualisation: is it needed to make visualisations
    """

    df = pd.read_csv(path_to_file)
    time_series = np.array(df['Level'])
    exog_variable = np.array(df['Neighboring level'])

    # Source time series
    train_input, predict_input = train_test_data_setup(InputData(idx=range(len(time_series)),
                                                                 features=time_series,
                                                                 target=time_series,
                                                                 task=Task(TaskTypesEnum.ts_forecasting,
                                                                           TsForecastingParams(
                                                                               forecast_length=len_forecast)),
                                                                 data_type=DataTypesEnum.ts))

    # Exogenous time series
    task = Task(TaskTypesEnum.ts_forecasting, TsForecastingParams(forecast_length=len_forecast))

    predict_input_exog = InputData(idx=np.arange(len(exog_variable)),
                                   features=exog_variable, target=time_series,
                                   task=task, data_type=DataTypesEnum.ts)

    train_input_exog, _ = train_test_data_setup(predict_input_exog)

    if with_exog is True:
        # Example with exogenous time series
        node_source = PrimaryNode('data_source_ts/1')
        node_lagged = SecondaryNode('lagged', nodes_from=[node_source])

        node_exog = PrimaryNode('exog_ts')

        node_final = SecondaryNode('ridge', nodes_from=[node_lagged, node_exog])
        pipeline = Pipeline(node_final)
    else:
        # Simple example without exogenous time series
        node_source_1 = PrimaryNode('data_source_ts/1')
        node_source_2 = PrimaryNode('data_source_ts/2')

        node_lagged_1 = SecondaryNode('lagged', nodes_from=[node_source_1])
        node_lagged_2 = SecondaryNode('lagged', nodes_from=[node_source_2])

        node_ridge_1 = SecondaryNode('ridge', nodes_from=[node_lagged_1])
        node_ridge_2 = SecondaryNode('ridge', nodes_from=[node_lagged_2])
        node_final = SecondaryNode('ridge', nodes_from=[node_ridge_1, node_ridge_2])
        train_input_exog = None
        predict_input_exog = None
        pipeline = Pipeline(node_final)

    predicted = make_forecast(pipeline, train_input, predict_input, train_input_exog, predict_input_exog)

    predicted = np.ravel(np.array(predicted))
    test_data = np.ravel(predict_input.target)

    print(f'Predicted values: {predicted[:5]}')
    print(f'Actual values: {test_data[:5]}')

    mse_before = mean_squared_error(test_data, predicted, squared=False)
    mae_before = mean_absolute_error(test_data, predicted)
    print(f'RMSE - {mse_before:.4f}')
    print(f'MAE - {mae_before:.4f}\n')

    if with_visualisation:
        plt.plot(range(0, len(time_series)), time_series, label='Actual time series')
        plt.plot(range(len(train_input.target), len(time_series)), predicted, label='Forecast')
        plt.legend()
        plt.grid()
        plt.show()

if __name__ == '__main__':
    data_path = os.path.join(f'{fedot_project_root()}', 'examples/data/ts', 'ts_sea_level.csv')
    run_exogenous_experiment(path_to_file=data_path, len_forecast=250, with_exog=True)

max_stars_repo_name: vishalbelsare/FEDOT

max_issues_repo_name: vishalbelsare/FEDOT

max_forks_repo_name: vishalbelsare/FEDOT


text:
# ******************************************************************************
# Copyright 2017-2018 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ******************************************************************************
from __future__ import print_function
import numpy as np
import pytest

import ngraph as ng
from ngraph.testing.flex_util import template_one_placeholder, MINIMUM_FLEX_VALUE, \
    MAXIMUM_FLEX_VALUE, id_func, template_two_placeholders, assert_allclose
from ngraph.testing import executor

pytestmark = pytest.mark.flex_only

# Known issues
bug_1064 = pytest.config.flex_disabled(reason="GitHub issue #1064, flex lower priority issues:"
                                              "modulus and ZeroDivisionError clarification")
bug_1227 = pytest.config.flex_disabled(reason="GitHub issue #1227, find explanation of results")

test_assign_data = (
    # template: ([(operand, expected_result, *case_description)], test_description),
    # *case_description is optional

    # test_assign
    ([(MINIMUM_FLEX_VALUE - 2, MINIMUM_FLEX_VALUE)], "Assign function - underflow expected"),
    ([(MAXIMUM_FLEX_VALUE + 1, MAXIMUM_FLEX_VALUE)], "Assign function - overflow expected"),
    ([(MINIMUM_FLEX_VALUE, MINIMUM_FLEX_VALUE)], "Assign function of negative boundary value"),
    ([(MAXIMUM_FLEX_VALUE, MAXIMUM_FLEX_VALUE)], "Assign function of positive boundary value"),
    ([(0.4, 0.399993896484375, "High precision check")],
     "Assign function of positive values from flex range - check high precision"),
    ([(1, 1),
      (5, 1.9999, "Operand overflow to 1.9999"),
      (10, 7.9997, "Operand overflow to 7.9997"),
      (MINIMUM_FLEX_VALUE, -32, "Operand underflow to -32")],
     "Iterations assign of x"),
)

test_data_single_operand = (
    # template: (ng_operation, [(operand, expected_result, *case_description)], test_description),
    # *case_description is optional

    # test_neg
    (ng.negative, [(MINIMUM_FLEX_VALUE, MAXIMUM_FLEX_VALUE)],
     "Negate function - overflow expected"),
    (ng.negative, [(MINIMUM_FLEX_VALUE + 1, MAXIMUM_FLEX_VALUE)],
     "Negate function of negative boundary value inside of flex range"),
    (ng.negative, [(MAXIMUM_FLEX_VALUE, MINIMUM_FLEX_VALUE + 1)],
     "Negate function of positive boundary value inside of flex range"),
    # test_sqrt
    (ng.sqrt, [(0, 0)], "Square root of zero equals zero"),
    (ng.sqrt, [(MAXIMUM_FLEX_VALUE, 181.015625, "High precision check")],
     "Square root of positive boundary value"),
    (ng.sqrt, [(MINIMUM_FLEX_VALUE, 0,
                "sqrt is implemented as LUT so input out of range saturates")],
     "Square root of negative boundary value"),
    (ng.sqrt, [(1, 1),
               (124, 1.4141, "Overflow of operand to 1.9999"),
               (10000, 2.8283, "Overflow of operand to 7.9997"),
               (32000, 5.6567, "Overflow of operand to 31.9990")],
     "Iterations sqrt of x "),

    # test_abs
    (ng.absolute, [(MINIMUM_FLEX_VALUE, MAXIMUM_FLEX_VALUE)],
     "Absolute value from the flex range - overflow expected"),
    (ng.absolute, [(MAXIMUM_FLEX_VALUE - 2, MAXIMUM_FLEX_VALUE - 2)],
     "Absolute value within of the flex range"),
    (ng.absolute, [(-1, 1),
                   (10000, 1.9999, "Overflow of operand to 1.9999"),
                   (-0.4, 0.3999),
                   (MINIMUM_FLEX_VALUE, 15.9995, "Overflow of operand to 15.9995")],
     "Iterations abs of x"),
    (ng.absolute, [(np.array([1, 2, 14, 4, 5, 6, 7, 8, 9, -1]),
                    np.array([1, 2, 14, 4, 5, 6, 7, 8, 9, 1]))],
     "Abs of array"),
    (ng.absolute, [(np.array([5, -1, 0, -2, 3, 4]).reshape(2, 3),
                    (np.array([5, 1, 0, 2, 3, 4]).reshape(2, 3)))], "Abs of 2x3 array"),
)

test_data_double_operand = (
    # template:(ng_operation, [(operand_1, operand_2, expected_result, *case_description)],
    # test_description), *case_description is optional

    # test_add
    (ng.add, [(MAXIMUM_FLEX_VALUE, 2, MAXIMUM_FLEX_VALUE)],
     "Positive boundary value plus one - overflow expected"),
    (ng.add, [(MINIMUM_FLEX_VALUE, 1, MINIMUM_FLEX_VALUE + 1)],
     "Negative boundary value plus one"),
    (ng.add, [(0, 1.5, 1.5),
              (1, 1.5, 1.99993896484375, "Overflow of result to 1.9999"),
              (2, 1.5, 3.5),
              (3, 1.5, 4.5)],
     "Iterations x + 1.5"),
    (ng.add, [(np.array([1, 1, 1, 1]), np.array([2, 3, 4, 5]), np.array([3, 4, 5, 6]))],
     "Add element-wise two arrays to each other"),

    # test_subtraction
    (ng.subtract, [(MINIMUM_FLEX_VALUE, 1, MINIMUM_FLEX_VALUE)],
     "Negative boundary value minus one - underflow expected"),
    (ng.subtract, [(MINIMUM_FLEX_VALUE, 2, MINIMUM_FLEX_VALUE)],
     "Negative boundary value minus two - underflow expected"),
    (ng.subtract, [(MAXIMUM_FLEX_VALUE, 1, MAXIMUM_FLEX_VALUE - 1)],
     "Positive boundary value minus one"),
    (ng.subtract, [(MAXIMUM_FLEX_VALUE, 2, MAXIMUM_FLEX_VALUE - 2)],
     "Positive boundary value minus two"),
    (ng.subtract, [(10, 0.4, 9.6),
                   (1000, 0.4, 15.5996, "Overflow of operand 1 to 15.9995"),
                   (10000, 0.4, 31.999, "Overflow of operand 1 to 63.9980"),
                   (1, 0.4, 0.6015)],
     "Iterations x - 0.4"),
    (ng.subtract, [(10000, 0.4, 9999.5, "Overflow of result to 9999.5"),
                   (5000, 0.4, 4999.5, "Overflow of result to 4999.5 "),
                   (2500, 0.4, 2500, "Overflow of result to 2500"),
                   (1000, 0.4, 1000, "Overflow of result to 1000"),
                   (800, 0.4, 800, "Overflow of result to 800"),
                   (750, 0.4, 750, "Overflow of result to 750"),
                   (600, 0.4, 600, "Overflow of result to 600"),
                   (500, 0.4, 500, "Overflow of result to 500"),
                   (300, 0.4, 300, "Overflow of result to 300"),
                   (200, 0.4, 200, "Overflow of result to 200"),
                   (250, 0.4, 250, "Overflow of result to 250"),
                   (100, 0.4, 100, "Overflow of result to 100"),
                   (80, 0.4, 80, "Overflow of result to 80"),
                   (75, 0.4, 75, "Overflow of result to 75"),
                   (60, 0.4, 60, "Overflow of result to 60"),
                   (50, 0.4, 50, "Overflow of result to 50"),
                   (30, 0.4, 30, "Overflow of result to 30"),
                   (25, 0.4, 24.5, "Overflow of result to 24.5"),
                   (20, 0.4, 19.5, "Overflow of result to 20"),
                   (10, 0.4, 9.625, "Overflow of result to 10"),
                   (1, 0.4, 0.625, "Overflow of result to 1")],
     "More complex iterations x - 0.4"),
    (ng.subtract, [(np.array([2, 3, 4, 5]), np.array([1, 1, 1, 1]), np.array([1, 2, 3, 4]))],
     "Subtract element-wise two arrays to each other"),

    # test_multiplication
    (ng.multiply, [(MINIMUM_FLEX_VALUE, 2, MINIMUM_FLEX_VALUE)],
     "Negative boundary value multiplied by two - underflow expected",),
    (ng.multiply, [(MAXIMUM_FLEX_VALUE, 2, MAXIMUM_FLEX_VALUE)],
     "Positive boundary value multiplied by two - overflow expected",),
    (ng.multiply, [(MINIMUM_FLEX_VALUE, 0, 0)],
     "Negative boundary value multiplied by zero equals zero"),
    (ng.multiply, [(MAXIMUM_FLEX_VALUE, 1, MAXIMUM_FLEX_VALUE)],
     "Positive boundary value multiplied by one is the same"),
    (ng.multiply, [(0, 0, 0)], "Zero multiplied by zero equals zero"),
    (ng.multiply, [(MINIMUM_FLEX_VALUE, -0.5, MINIMUM_FLEX_VALUE * (-0.5))],
     "Negative sign value multiplied by negative sign value equals positive sign"),
    (ng.multiply, [(1, 10.1, 10.0996),
                   (1000, 10.1, 15.9995, "Overflow of operand 1 to 1.9999 and result to 15.999"),
                   (2, 10.1, 20.1992,),
                   (-1000, 10.1, -64, "Underflow of operand 1 to -8 and result to -64"),
                   (0.4, 10.1, 4.0312)],
     "Iterations x * 10.1"),
    (ng.multiply, [(np.array([2, 3, 4, 5]), np.array([1, 1, 1, 1]), np.array([2, 3, 4, 5]))],
     "Multiply element-wise two arrays to each other"),

    # test_division
    (ng.divide, [(MAXIMUM_FLEX_VALUE, 0.5, MAXIMUM_FLEX_VALUE)],
     "Positive boundary value division - overflow expected"),
    (ng.divide, [(MINIMUM_FLEX_VALUE, 0.5, MINIMUM_FLEX_VALUE)],
     "Negative boundary value division - underflow expected"),
    bug_1227((ng.divide, [(MAXIMUM_FLEX_VALUE, 3, 10922)], "Positive boundary value division")),
    bug_1227((ng.divide, [(MINIMUM_FLEX_VALUE, 3, -10922)], "Negative boundary value division")),
    (ng.divide, [(-10, 7, -1.4285),
                 (0.4, 7, 0.0571),
                 (MAXIMUM_FLEX_VALUE, 7, 3.9998,
                  "Overflow operand 1 to 31.9990 and result to 3,9998"),
                 (MINIMUM_FLEX_VALUE, 7, -16, "Underflow operand 1 to -128 and result to -16"),
                 (0, 7, 0)],
     "Iterations x / 7"),

    # test_modulo
    bug_1064((ng.mod, [(MINIMUM_FLEX_VALUE, 3, MINIMUM_FLEX_VALUE % 3)],
              "Negative boundary value mod 3")),
    (ng.mod, [(MAXIMUM_FLEX_VALUE, 3, MAXIMUM_FLEX_VALUE % 3)], "Positive boundary value mod 3"),
    bug_1064((ng.mod, [(MAXIMUM_FLEX_VALUE, -3, MAXIMUM_FLEX_VALUE % (-3))],
              "Positive boundary value mod -3")),
    bug_1064((ng.mod, [(2.1, 2, 0.09999847412109375, "High precision check")],
              "Modulo of floating point")),

    # test_power
    (ng.power, [(MAXIMUM_FLEX_VALUE, 2, MAXIMUM_FLEX_VALUE)],
     "Positive boundary value exponentiation - overflow expected"),
    (ng.power, [(MINIMUM_FLEX_VALUE, 3, MINIMUM_FLEX_VALUE)],
     "Negative boundary value exponentiation - underflow expected"),
    # Not sure of this case, results should be tracked
    (ng.power, [(MAXIMUM_FLEX_VALUE, 0.4, 63.99609375, "High precision check")],
     "Positive boundary value exponentiation"),
    (ng.power, [(MINIMUM_FLEX_VALUE, -2, MINIMUM_FLEX_VALUE ** (-2))],
     "Negative boundary value negative exponentiation"),
    (ng.power, [(1, 3, 1),
                (5, 3, 1.9999, "Underflow operand 1 and result to 1.9999"),
                (-10, 3, -8, "Underflow operand 1 and result to -8"),
                (22, 3, 31.999, "Overflow result to 31.999")],
     "Iterations x ^ 3")
)

@pytest.mark.parametrize("operation, operands, test_name", test_data_single_operand, ids=id_func)
def test_single_operand(operation, operands, test_name):
    template_one_placeholder(operands, operation)

@pytest.mark.parametrize("operation, operands, test_name", test_data_double_operand, ids=id_func)
def test_double_operand(operation, operands, test_name):
    template_two_placeholders(operands, operation)

@pytest.mark.parametrize("operands, test_name", test_assign_data, ids=id_func)
def test_assign(operands, test_name):
    v = ng.variable(())
    ng_placeholder = ng.placeholder(())
    vset = ng.sequential([
        ng.assign(v, ng_placeholder),
        v
    ])
    iterations = len(operands) != 1
    with executor(vset, ng_placeholder) as ex:
        for i in operands:
            flex_result = ex(i[0])
            print("flex: ", flex_result)
            print("expected: ", i[1])
            if iterations:
                assert_allclose(flex_result, i[1])
            else:
                assert flex_result == i[1]

max_stars_repo_name: NervanaSystems/ngraph-python


text:
import numpy as np

import scipy.spatial as sspatial
from sklearn.neighbors import KRadiusNeighborsRegressor, NearestNeighbors

import optim
import poli

from GPy.models import GPRegression

import matplotlib.pyplot as plt
import matplotlib.cm as mcm

class NNR(object):
    def __init__(self, r=1.0, k=10, def_mean=0, def_sd=float('inf'), off=0.1):
        self.reg = KRadiusNeighborsRegressor(
            n_neighbors=k, radius=r, defval=def_mean, weights=self.comp_weights)
        # Regresses squared error
        self.err_reg = NearestNeighbors(n_neighbors=k, radius=r)
        self.off = off
        self.def_sd = float(def_sd)

    def comp_weights(self, dists):
        return 1.0 / (dists + self.off)

    def fit(self, X, y):
        self.reg.fit(X, y)
        errs = y - self.reg.predict(X)
        self.sse = errs * errs
        self.err_reg.fit(X)

    def predict(self, X, return_std=False):
        if not return_std:
            return self.reg.predict(X)
        else:
            mean_val = self.reg.predict(X)

            sds = []
            dists, inds = self.err_reg.radius_neighbors(
                X, return_distance=True)
            for d, i in zip(dists, inds):
                if len(d) < 2:
                    sds.append(self.def_sd)
                    continue
                else:
                    errs = self.sse[i]
                    weights = self.comp_weights(d)
                    sd = np.average(errs, axis=0, weights=weights) / len(d)
                    sds.append(sd)

            return mean_val, np.array(sds)

class TestModalProblem(object):
    def __init__(self, x_dim, a_dim, n_modes):
        self.centers = np.random.uniform(low=-1, high=1, size=(n_modes, x_dim))
        self.offsets = np.random.uniform(low=-1, high=1, size=(n_modes, a_dim))

        self.current_context = np.random.uniform(low=-1, high=1, size=x_dim)

    @property
    def context(self):
        return self.current_context

    def step(self, a):
        x = self.current_context.reshape(1, -1)
        dists = sspatial.distance.cdist(x, self.centers)
        # probs = np.exp( -dists )
        # probs = probs / np.sum(probs)
        # ind = np.random.choice(range(len(dists)), p=probs)
        ind = np.argmin(dists)

        noise = np.random.normal(scale=0.1)
        reward = -np.log(np.linalg.norm(self.offsets[ind] - a)) + noise
        reward = max(min(reward, 5), -5)

        self.current_context = np.random.uniform(low=-1, high=1, size=x_dim)
        return reward

if __name__ == '__main__':

    x_dim = 2
    a_dim = 1
    n_modes = 3

    problem = TestModalProblem(x_dim=x_dim, a_dim=a_dim, n_modes=n_modes)
    mp = poli.ModalPolicy()

    n_init = 30
    n_opt_iters = 5
    n_trials = 1000
    opt_per = 10
    rewards = []
    states_seen = []

    for i in range(n_init):
        s = problem.context
        a = np.random.uniform(-1, 1, a_dim)
        r = problem.step(a)
        states_seen.append(s)
        rewards.append(r)

        print 's: %s a: %s r: %f' % (np.array_str(s), np.array_str(a), r)

        mp.report_sample(s=s, a=a, r=r)

    def print_cluster():
        print 'clusters:\n'
        for cc, r in zip(mp.cluster_actions, mp.cluster_rewards):
            print '\tc: %s r: %f' % (np.array_str(cc), r)

    mp.initialize_modes(n_modes * 2)
    mp.optimize(n_opt_iters, beta=1)
    print_cluster()

    # plotting
    grid = np.linspace(-1, 1, 100)
    xx, yy = np.meshgrid(grid.flatten(), grid.flatten())
    X = np.vstack((xx.flatten(), yy.flatten())).T

    plt.ion()
    plt.figure()

    def update_plot():
        zz = mp.clusters.predict(X)
        Z = np.reshape(zz, xx.shape)
        plt.gca().clear()
        plt.gca().set_xlim((-1,1))
        plt.gca().set_ylim((-1,1))
        plt.pcolormesh(xx, yy, Z, cmap=mcm.PiYG)
        #plt.contourf(xx, yy, Z, np.unique(zz), antialiased=False)
        plt.plot(problem.centers[:,0], problem.centers[:,1], 'kx', markersize=5, mew=3)
        s = np.array(states_seen)
        plt.scatter(x=s[:,0], y=s[:,1], marker='.', color='k')
        plt.draw()

    for i in range(n_trials):
        s = problem.context
        a = mp.get_action(s)
        r = problem.step(a)
        rewards.append(r)
        states_seen.append(s)

        print 's: %s a: %s r: %f' % (np.array_str(s), np.array_str(a), r)

        # if i % opt_per == 0:
        mp.report_sample(s=s, a=a, r=r)
        mp.optimize(1, beta=1)
        print_cluster()
        update_plot()

    plt.figure('Reward trace')
    plt.title('Rewards vs time')
    plt.plot(rewards, 'r-')
    plt.xlabel('Iteration')
    plt.ylabel('Reward')

max_stars_repo_name: Humhu/percepto

max_issues_repo_name: Humhu/percepto

max_forks_repo_name: Humhu/percepto

reason: import numpy,import scipy