text:
#!/usr/bin/env python-sirius

"""."""

import os as _os
import sys as _sys
from siriuspy import util as _util
import pickle as _pic
import numpy as _np

# TODO: funcionalities of this module are being moved to
# fieldmaptrack.hallsensor eventually this module should be deleted and thus
# it should not be maintained or used

class Templates:
    """."""

    rawfield = """
    # ==========================
    # fma_rawfield.py input file
    # Date: TIMESTAMP
    # Accelerator Physics LNLS
    # ==========================

    # --- Summary ---
    #
    # this is the input file for fac-fma-rawfield.py script
    # this script reads a fieldmap from a 3D magnet model, stores it
    # for latter analysis and prints and plots basic information on the
    # field map. It is used to quickly inspect the fieldmap

    # --- Input parameters ---

    # each analysis has an identity label used for naming output files

      config_label             'CONFIG_NAME'

    # the next parameter specifies the type of magnet to be analysed.
    # each type may have its own particular algorithms to be applied

      magnet_type              'dipole'

    # the full name of the file that contains the field map

      fmap_filename            'FIELDMAP_FNAME'

    # Runge-Kutta algorithm used for the integration of the eqs. of motion needs to know
    # what to do when trajectory reaches the fieldmap bounds. It will either extrapolate the field
    # along the longitudinal (z) direction or consider it to have vanished. This is controlled with
    # the parameter below. Bear in mind that the calculation of extrapolation coefficients is very
    # time-consuming currently. As for the transverse directions (x and y), the RK algorithm will
    # generate exceptions.

      fmap_extrapolation_flag  False

    """

    trajectory = """
    # ==========================
    # fma_rawfield.py input file
    # Date: TIMESTAMP
    # Accelerator Physics LNLS
    # ==========================

    # --- Summary ---
    #
    # This is the input file for trajectory calculation based on a given
    # fieldmap which is performed with the script 'fac-fma-trajectory.py'
    # A controllable fixed-size Runge-Kutta algorithm is used to integrate
    # the equations of motion of a single electron in the presence of
    # the magnetic field as defined in the fieldmap.
    #
    # The implemented equations of motion are not approximated. Provided
    # a sufficiently fine RK step is chosen, this scripts may be used to
    # accurately obtain the trajectory of the electron with arbitrary energy
    #
    # Runge-Kutta algorithm used for the integration of the eqs. of motion needs to know
    # what to do when trajectory reaches the fieldmap bounds. It will either extrapolate the field
    # along the longitudinal (z) direction or consider it to have vanished.
    # As for the transverse directions (x and y), the RK algorithm will
    # generate exceptions.

    # --- Input parameters ---

    # each analysis has an identity label used for naming output files

      config_label             	'CONFIG_NAME'

    # beam energy

      beam_energy                     BEAM_ENERGY     # [GeV]

    # A trajectory can also be read from file. This is useful when the fieldmap of
    # 3D models with errors are being analysed. In this case we want to use as reference
    # trajectory a trajectory that was calculated from the 3D model without errors and
    # saved to file. If parameter 'traj_load_filename' is set to 'None' then a new
    # reference trajectory with be calculated with RK on the given fieldmap.

      traj_load_filename              None

    # If parameter 'traj_is_reference_traj' is set to True the algorithm will rescale the
    # fieldmap so that the total trajectory deflection will exactly match the nominal deflection

      traj_is_reference_traj          False
      model_nominal_angle             7.2                           # [deg] model nominal deflection angle of the magnet

    # There is the option to restrain the trajectory to the midplane (y = 0 mm) of the magnet

      traj_force_midplane_flag        True

    # There is the option to serach for a initial rx position that will result in a trajectory that
    # is centered in the good-field region of the magnet (around rx == 0)

      traj_center_sagitta_flag        False

    # The RK algorithm always integrates the trajectory from the center of the magnet (z = s = 0 mm)
    # The limits of the RK integration may be specified in various ways:
    # If only 'traj_rk_s_step' is given then the algorithm will integrate until
    # the z coordinate of the particle reaches the fieldmap bound.

      traj_init_rx                    RX_INIT                       # [mm] init_rx at z = s = 0 mm (center of magnet)
      traj_rk_s_step                  S_STEP                        # [mm]
      traj_rk_length                  None                          # [mm]
      traj_rk_nrpts                   None

    # whether to save trajectory to an ASCII file

      traj_save                       True
    """

    multipoles = """
    # ==========================
    # fma_rawfield.py input file
    # Date: TIMESTAMP
    # Accelerator Physics LNLS
    # ==========================

    # --- Summary ---
    #
    # this is the input file for the 'fac-fma-multipoles.py' script
    # this script calculates the multipoles around the reference trajectory.

    # --- Input parameters ---

    # each analysis has an identity label used for naming output files

      config_label                      'CONFIG_NAME'

    # the multipoles (m1,m2,...) to be calculated are defined by a list of position x exponents (n1,n2,...):
    # By = m1 * x^n1 + m2 * x^n2 + ...

      multipoles_normal_field_fitting_monomials      (0,1,2,3,4,5,6)                 # monomials to be included in the polynomial fit of multipoles
      multipoles_skew_field_fitting_monomials        ()

    # grid of perpendicular points around each point of the reference trajectory for the polynomial fit of By and Bx

      multipoles_perpendicular_grid     np.linspace(-12,12,65)          # grid of points on perpendicular line to ref trajectory [mm]

    # after multipole coeffs are calculated, their normalized strengths at perp. position r0 are calculated (as defined in tracy)

      multipoles_r0                     17.5                             # [mm] horizontal position at which polynomial fields are calculated relative to the principal multipole
      normalization_monomial            0
      normalization_is_skew             False

    # integrated residual field (converted to kick angle) calculated from fitted multipoles and
    # from integrated fieldmap are compared. The parameter below lists the monomials which are
    # supposed to define the main field. The rest makes up for the residual field

      normal_multipoles_main_monomials         (0,1,2)
      skew_multipoles_main_monomials           ()
    """

    model = """
    # ==========================
    # fma_rawfield.py input file
    # Date: 2018-07-20
    # Accelerator Physics LNLS
    # ==========================

    # --- Summary ---
    #
    # This script integrates fitted multipoles at each segment of the hard-edge model

    # --- Input parameters ---

    # each analysis has an identity label used for naming output files

      config_label                      'CONFIG_NAME'

    # list with lengths of model segments

       model_segmentation               (196, 192, 182, 10, 10, 13, 17, 20, 30, 50)
    """

    help = """
    NAME
           hallprobe.py - routines and libs to process fieldmap analysis

    SYNOPSIS
           hallprobe.py [CMD] [ARGS]

    DESCRIPTION
           The command options are

           help
           summary  [x0-9p1013mm|...] [current_label] [positive|negative]
           new-energy [x0-9p1013mm|...]
    """

# text templates
_rawfield = Templates.rawfield
_trajectory = Templates.trajectory
_multipoles = Templates.multipoles
_model = Templates.model
_help = Templates.help

class FMapAnalysisTB:
    """."""

    def __init__(self,
                 magnet,
                 curlabel,
                 path_analysis,
                 path_fmap,
                 beam_energy,
                 rx_init,
                 s_step):
        """."""
        self.magnet = magnet
        self.curlabel = curlabel
        self.path_analysis = path_analysis
        self.path_fmap = path_fmap
        self.beam_energy = beam_energy
        self.rx_init = rx_init
        self.s_step = s_step
        if not _os.path.exists(path_analysis):
            _os.makedirs(path_analysis)

    def files_create(self):
        """."""
        self._create_rawfield_in()
        self._create_trajectory_in()
        self._create_multipoles_in()
        self._create_model_in()

    def files_read_results_rawfield(self):
        """."""
        fname = self.path_analysis + '/rawfield.out'
        with open(fname, 'r') as f:
            lines = f.readlines()
        for line in lines:
            if 'main_coil_current' in line:
                self.results_current = float(line.split()[1])

    def files_read_results_trajectory(self):
        """."""
        fname = self.path_analysis + '/trajectory.out'
        with open(fname, 'r') as f:
            lines = f.readlines()
        for line in lines:
            if 'horizontal_deflection_angle' in line:
                self.results_angle = float(line.split()[1])
            elif 'rx position of reference' in line:
                self.results_rx_ref = float(line.split()[5])
            # elif 'initial rx position' in line:
            #     self.results_rx_init = float(line.split()[5])

    def files_read_results_multipoles(self):
        """."""
        fname = self.path_analysis + '/multipoles.out'
        with open(fname, 'r') as f:
            lines = f.readlines()

        self.multipoles_normal = {}
        for line in lines:
            if 'r0_for_relative_multipoles' in line:
                self.reference_r0 = float(line.split()[1])/1000.0
            for n in range(30):
                nstr = 'n={:02d}'.format(n)
                if nstr in line:
                    dstr = line.split()
                    self.multipoles_normal[n] = float(dstr[2])

    def cmd_clean(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py clean'
        _os.system(cmd)

    def cmd_rawfield(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py rawfield'
        _os.system(cmd)

    def cmd_trajectory(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py trajectory'
        _os.system(cmd)

    def cmd_multipoles(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py multipoles'
        _os.system(cmd)

    def cmd_model(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py model'
        _os.system(cmd)

    def cmd_run(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py run'
        _os.system(cmd)

    def analysis_trajectory(self):
        """."""
        self.cmd_clean()
        self.files_create()
        self.cmd_rawfield()
        self.cmd_trajectory()
        self.files_read_results_rawfield()
        self.files_read_results_trajectory()

    # --- auxiliary method ---

    def _create_rawfield_in(self):
        text = self._process_text(_rawfield)
        fname = self.path_analysis + '/rawfield.in'
        with open(fname, 'w') as f:
            f.write(text)

    def _create_trajectory_in(self):
        text = self._process_text(_trajectory)
        fname = self.path_analysis + '/trajectory.in'
        with open(fname, 'w') as f:
            f.write(text)

    def _create_multipoles_in(self):
        text = self._process_text(_multipoles)
        fname = self.path_analysis + '/multipoles.in'
        with open(fname, 'w') as f:
            f.write(text)

    def _create_model_in(self):
        text = self._process_text(_model)
        fname = self.path_analysis + '/model.in'
        with open(fname, 'w') as f:
            f.write(text)

    def _process_text(self, text):
        t = text
        t = t.replace('TIMESTAMP', _util.get_timestamp())
        t = t.replace('MAGNET', self.magnet)
        t = t.replace('CONFIG_NAME', self.magnet+'-'+self.curlabel)
        t = t.replace('FIELDMAP_FNAME', self.path_fmap)
        t = t.replace('BEAM_ENERGY', str(self.beam_energy))
        t = t.replace('RX_INIT', str(self.rx_init))
        t = t.replace('S_STEP', str(self.s_step))
        return t

class FMapAnalysisBO:
    """."""

    def __init__(self,
                 magnet,
                 curlabel,
                 path_analysis,
                 path_fmap,
                 beam_energy,
                 rx_init,
                 s_step):
        """."""
        self.magnet = magnet
        self.curlabel = curlabel
        self.path_analysis = path_analysis
        self.path_fmap = path_fmap
        self.beam_energy = beam_energy
        self.rx_init = rx_init
        self.s_step = s_step
        if not _os.path.exists(path_analysis):
            _os.makedirs(path_analysis)

    def files_create(self):
        """."""
        self._create_rawfield_in()
        self._create_trajectory_in()
        self._create_multipoles_in()
        self._create_model_in()

    def files_read_results_rawfield(self):
        """."""
        fname = self.path_analysis + '/rawfield.out'
        with open(fname, 'r') as f:
            lines = f.readlines()
        for line in lines:
            if 'main_coil_current' in line:
                self.results_current = float(line.split()[1])

    def files_read_results_trajectory(self):
        """."""
        fname = self.path_analysis + '/trajectory.out'
        with open(fname, 'r') as f:
            lines = f.readlines()
        for line in lines:
            if 'horizontal_deflection_angle' in line:
                self.results_angle = float(line.split()[1])
            elif 'rx position of reference' in line:
                self.results_rx_ref = float(line.split()[5])
            elif 'initial rx position' in line:
                self.results_rx_init = float(line.split()[5])

    def files_read_results_multipoles(self):
        """."""
        fname = self.path_analysis + '/multipoles.out'
        with open(fname, 'r') as f:
            lines = f.readlines()
        for line in lines:
            if 'n=00' in line:
                self.results_intn0 = float(line.split()[2])
            elif 'n=01' in line:
                self.results_intn1 = float(line.split()[2])
            elif 'n=02' in line:
                self.results_intn2 = float(line.split()[2])

    def cmd_clean(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py clean'
        _os.system(cmd)

    def cmd_rawfield(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py rawfield'
        _os.system(cmd)

    def cmd_trajectory(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py trajectory'
        _os.system(cmd)

    def cmd_multipoles(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py multipoles'
        _os.system(cmd)

    def cmd_model(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py model'
        _os.system(cmd)

    def cmd_run(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; fac-fma-analysis.py run'
        _os.system(cmd)

    def analysis_trajectory(self):
        """."""
        self.cmd_clean()
        self.files_create()
        self.cmd_rawfield()
        self.cmd_trajectory()
        self.files_read_results_rawfield()
        self.files_read_results_trajectory()

    def analysis_multipoles(self):
        """."""
        self.cmd_clean()
        self.files_create()
        self.cmd_rawfield()
        self.cmd_trajectory()
        self.cmd_multipoles()
        self.files_read_results_rawfield()
        self.files_read_results_trajectory()
        self.files_read_results_multipoles()

    def analysis_model(self):
        """."""
        self.cmd_clean()
        self.files_create()
        self.cmd_rawfield()
        self.cmd_trajectory()
        self.cmd_multipoles()
        self.cmd_model()
        self.files_read_results_rawfield()
        self.files_read_results_trajectory()
        self.files_read_results_multipoles()

    def analysis_concatenate_output_files(self):
        """."""
        path = self.path_analysis
        cmd = 'cd ' + path + '; cat rawfield.out >> analysis.txt'
        _os.system(cmd)
        cmd = 'cd ' + path + '; cat trajectory.out >> analysis.txt'
        _os.system(cmd)
        cmd = 'cd ' + path + '; cat multipoles.out >> analysis.txt'
        _os.system(cmd)
        cmd = 'cd ' + path + '; cat model.out >> analysis.txt'
        _os.system(cmd)

    # --- auxiliary method ---

    def _create_rawfield_in(self):
        text = self._process_text(_rawfield)
        fname = self.path_analysis + '/rawfield.in'
        with open(fname, 'w') as f:
            f.write(text)

    def _create_trajectory_in(self):
        text = self._process_text(_trajectory)
        fname = self.path_analysis + '/trajectory.in'
        with open(fname, 'w') as f:
            f.write(text)

    def _create_multipoles_in(self):
        text = self._process_text(_multipoles)
        fname = self.path_analysis + '/multipoles.in'
        with open(fname, 'w') as f:
            f.write(text)

    def _create_model_in(self):
        text = self._process_text(_model)
        fname = self.path_analysis + '/model.in'
        with open(fname, 'w') as f:
            f.write(text)

    def _process_text(self, text):
        t = text
        t = t.replace('TIMESTAMP', _util.get_timestamp())
        t = t.replace('MAGNET', self.magnet)
        t = t.replace('CONFIG_NAME', self.magnet+'-'+self.curlabel)
        t = t.replace('FIELDMAP_FNAME', self.path_fmap)
        t = t.replace('BEAM_ENERGY', str(self.beam_energy))
        t = t.replace('RX_INIT', str(self.rx_init))
        t = t.replace('S_STEP', str(self.s_step))
        return t

class FMapAnalysisProductionTB:
    """."""

    _path_analysis = (
        '/home/fac_files/lnls-ima/tb-dipoles/model-03'
        '/analysis/hallprobe/excitation_curve/x0-4p6467mm-15deg/')

    _path_measurements = (
        '/home/fac_files/lnls-ima/tb-dipoles/model-03'
        '/measurement/magnetic/hallprobe/excitation_curve/')

    def __init__(self,
                 magnet,
                 curlabel,
                 fmap_fname,
                 beam_energy,
                 rx_init,
                 s_step):
        """."""
        self.magnet = magnet
        self.curlabel = curlabel
        self.beam_energy = beam_energy
        self.path_analysis = FMapAnalysisProductionTB._path_analysis + \
            self.magnet + '/' + self.curlabel + '/'
        self.path_measurement = FMapAnalysisProductionTB._path_measurements + \
            self.magnet + '/'
        self.path_fmap = self.path_measurement + fmap_fname
        self.s_step = abs(s_step)
        self.analysis_neg = FMapAnalysisTB(self.magnet,
                                           self.curlabel,
                                           self.path_analysis + 'z-negative',
                                           self.path_fmap,
                                           self.beam_energy,
                                           rx_init,
                                           -1.0*self.s_step)
        self.analysis_pos = FMapAnalysisTB(self.magnet,
                                           self.curlabel,
                                           self.path_analysis + 'z-positive',
                                           self.path_fmap,
                                           self.beam_energy,
                                           rx_init,
                                           +1.0*self.s_step)

    def files_create(self):
        """."""
        self.analysis_pos.files_create()
        self.analysis_neg.files_create()

    def files_clean(self):
        """."""
        self.analysis_pos.files_clean()
        self.analysis_neg.files_clean()

    def files_read_results(self):
        """."""
        self.analysis_pos.files_read_results_rawfield()
        self.analysis_neg.files_read_results_rawfield()
        self.analysis_pos.files_read_results_trajectory()
        self.analysis_neg.files_read_results_trajectory()
        self.results_angle = self.analysis_pos.results_angle - \
            self.analysis_neg.results_angle
        self.analysis_pos.files_read_results_multipoles()
        self.analysis_neg.files_read_results_multipoles()
        self.multipoles_normal = {}
        self.multipoles_normal_relative = {}
        r0 = self.analysis_pos.reference_r0
        mm = (self.analysis_pos.multipoles_normal[0] +
              self.analysis_neg.multipoles_normal[0]) * r0**0
        for n, v in self.analysis_pos.multipoles_normal.items():
            self.multipoles_normal[n] = \
                self.analysis_pos.multipoles_normal[n] + \
                self.analysis_neg.multipoles_normal[n]
            self.multipoles_normal_relative[n] = \
                self.multipoles_normal[n] * (r0**n) / mm

    def run(self):
        """."""
        self.analysis_pos.files_run()
        self.analysis_neg.files_run()

    def calc_multipoles_kick(self, xmax,
                             multipoles_normal=None,
                             excluded_monomials=None):
        """."""
        if excluded_monomials is None:
            excluded_monomials = (0,)
        if multipoles_normal is None:
            multipoles_normal = self.multipoles_normal
        x = _np.linspace(-xmax, xmax, 101)
        y = _np.zeros(x.shape)
        for n, v in multipoles_normal.items():
            if n not in excluded_monomials:
                y += v * x**n
        brho, *_ = _util.beam_rigidity(self.beam_energy)
        y /= brho
        return y, x

class FMapAnalysisProductionBO:
    """."""

    _path_analysis = (
        '/home/fac_files/lnls-ima/bo-dipoles/model-09'
        '/analysis/hallprobe/production/')

    _path_measurements = (
        '/home/fac_files/lnls-ima/bo-dipoles/model-09'
        '/measurement/magnetic/hallprobe/production/')

    def __init__(self,
                 magnet,
                 curlabel,
                 fmap_fname,
                 beam_energy,
                 s_step):
        """."""
        self.magnet = magnet
        self.curlabel = curlabel
        self.beam_energy = beam_energy
        self.path_analysis = FMapAnalysisProductionBO._path_analysis + \
            self.magnet + '/' + 'M1' + '/' + self.curlabel + '/'
        self.path_measurement = FMapAnalysisProductionBO._path_measurements + \
            self.magnet + '/' + 'M1' + '/'
        self.path_fmap = self.path_measurement + fmap_fname
        self.s_step = abs(s_step)
        self.analysis_neg = FMapAnalysisBO(self.magnet,
                                           self.curlabel,
                                           self.path_analysis + 'z-negative',
                                           self.path_fmap,
                                           self.beam_energy,
                                           -1.0*self.s_step)
        self.analysis_pos = FMapAnalysisBO(self.magnet,
                                           self.curlabel,
                                           self.path_analysis + 'z-positive',
                                           self.path_fmap,
                                           self.beam_energy,
                                           +1.0*self.s_step)

    def files_create(self):
        """."""
        self.analysis_pos.files_create()
        self.analysis_neg.files_create()

    def files_clean(self):
        """."""
        self.analysis_pos.files_clean()
        self.analysis_neg.files_clean()

    def run(self):
        """."""
        self.analysis_pos.files_run()
        self.analysis_neg.files_run()

class HallProbeAnalysisBO:
    """."""

    _path = ('/home/fac_files/lnls-ima/bo-dipoles/model-09/'
             'analysis/hallprobe/excitation_curve/')

    def __init__(self, dataset, name, current):
        """."""
        self.dataset = dataset
        self.name = name
        self.current = current
        self.current_pos, self.beam_energy_pos, \
            self.angle_pos, \
            self.intn0_pos, self.intn1_pos, self.intn2_pos = \
            self._read_analysis('z-positive')
        self.current_neg, self.beam_energy_neg, \
            self.angle_neg, \
            self.intn0_neg, self.intn1_neg, self.intn2_neg = \
            self._read_analysis('z-negative')
        self.current = 0.5*(self.current_pos + self.current_neg)
        self.beam_energy = 0.5*(self.beam_energy_pos + self.beam_energy_neg)
        self.angle = self.angle_pos - self.angle_neg  # corrects sign error
        self.intn0 = self.intn0_pos + self.intn0_neg
        self.intn1 = self.intn1_pos + self.intn1_neg
        self.intn2 = self.intn2_pos + self.intn2_neg

    def _read_analysis(self, side):

        fname = HallProbeAnalysisBO._path + self.dataset + '/' + \
                self.name + '/' + self.current + '/' + side + '/analysis.txt'
        with open(fname, 'r') as f:
            lines = f.readlines()
        for line in lines:
            if 'horizontal_deflection_angle' in line:
                angle = float(line.split()[1])
            elif 'n=00' in line:
                intn0 = float(line.split()[2])
            elif 'n=01' in line:
                intn1 = float(line.split()[2])
            elif 'n=02' in line:
                intn2 = float(line.split()[2])
            elif 'main_coil_current' in line:
                current = float(line.split()[1])
            elif 'beam_energy' in line:
                energy = float(line.split()[1])
        return current, energy, angle, intn0, intn1, intn2

class DipolesSetBO:
    """."""

    spec_angle = -7.2  # [deg]
    spec_angle_rms = 0.15  # [%]
    spec_ref_current = 991.63  # [A]
    spec_quadrupole = 2.4788148246867  # [T] @991.63A (3 GeV)
    spec_quadrupole_rms = 2.4  # [%]
    spec_sextupole = 25.627729062301  # [T/m] @991.63A (3 GeV)
    spec_sextupole_rms = 9.0  # [%]

    def __init__(self, dataset, magnets):
        """."""
        self.dataset = dataset
        self._data = dict()
        for magnet, currlabels in magnets.items():
            self._data[magnet] = dict()
            for currlabel in currlabels:
                self._data[magnet][currlabel] = \
                    HallProbeAnalysisBO(self.dataset, magnet, currlabel)

    @property
    def magnet_labels(self):
        """."""
        return list(self._data.keys())

    def currents_minmax(self):
        """."""
        mi, ma = None, None
        for v in self._data.values():
            currlabels = list(v.keys())
            currents = [v[c].current for c in currlabels]
            mi = min(currents) if not mi else min(mi, min(currents))
            ma = max(currents) if not ma else max(ma, max(currents))
        return mi, ma

    def get_energies(self):
        """."""
        maglabel_set = self.magnet_labels
        currents_set = []
        energies_set = []
        for v in self._data.values():
            currlabels = list(v.keys())
            currents = [v[c].current for c in currlabels]
            energies = [v[c].beam_energy for c in currlabels]
            currents_set.append(currents)
            energies_set.append(energies)
        return maglabel_set, currents_set, energies_set

    def get_angles(self):
        """."""
        maglabel_set = self.magnet_labels
        currents_set = []
        angles_set = []
        for v in self._data.values():
            currlabels = list(v.keys())
            currents = [v[c].current for c in currlabels]
            angles = [v[c].angle for c in currlabels]
            currents_set.append(currents)
            angles_set.append(angles)
        return maglabel_set, currents_set, angles_set

    def get_quadrupoles(self):
        """."""
        maglabel_set = self.magnet_labels
        currents_set = []
        quadrupoles_set = []
        for v in self._data.values():
            currlabels = list(v.keys())
            currents = [v[c].current for c in currlabels]
            quadrupoles = [v[c].intn1 for c in currlabels]
            currents_set.append(currents)
            quadrupoles_set.append(quadrupoles)
        return maglabel_set, currents_set, quadrupoles_set

    def get_sextupoles(self):
        """."""
        maglabel_set = self.magnet_labels
        currents_set = []
        sextupoles_set = []
        for v in self._data.values():
            currlabels = list(v.keys())
            currents = [v[c].current for c in currlabels]
            sextupoles = [v[c].intn2 for c in currlabels]
            currents_set.append(currents)
            sextupoles_set.append(sextupoles)
        return maglabel_set, currents_set, sextupoles_set

    def plot_energies(self, plt):
        """."""
        magnets, currents, energies = self.get_energies()
        for i in range(len(magnets)):
            plt.plot(currents[i], energies[i], 'o-', label=magnets[i])
        plt.xlabel('Current [A]')
        plt.ylabel('Energy [GeV]')
        plt.legend()

    def plot_angles(self, plt):
        """."""
        magnets, currents, angles = self.get_angles()
        spec = DipolesSetBO.spec_angle
        xrms = self.currents_minmax()
        plt.plot(xrms, (+DipolesSetBO.spec_angle_rms,)*2, 'k--')
        plt.plot(xrms, (-DipolesSetBO.spec_angle_rms,)*2, 'k--')
        for i in range(len(magnets)):
            error = 100 * (_np.array(angles[i]) - spec)/spec
            plt.plot(currents[i], error, 'o-', label=magnets[i])
        plt.xlabel('Current [A]')
        plt.ylabel('Angle Error w.r.t. to Spec [%]')
        plt.legend()

    def plot_quadrupoles(self, plt):
        """."""
        magnets, currents, quadrupoles = self.get_quadrupoles()
        spec0 = DipolesSetBO.spec_quadrupole
        xrms = self.currents_minmax()
        plt.plot(xrms, (+DipolesSetBO.spec_quadrupole_rms,)*2, 'k--')
        plt.plot(xrms, (-DipolesSetBO.spec_quadrupole_rms,)*2, 'k--')
        for i in range(len(magnets)):
            error = []
            for j in range(len(currents[i])):
                spec = spec0 * currents[i][j]/DipolesSetBO.spec_ref_current
                error.append(100 * (quadrupoles[i][j] - spec)/spec)
            plt.plot(currents[i], error, 'o-', label=magnets[i])
        plt.xlabel('Current [A]')
        plt.ylabel('Quadrupole Error w.r.t. to Spec [%]')
        plt.legend()

    def plot_sextupoles(self, plt):
        """."""
        magnets, currents, sextupoles = self.get_sextupoles()
        spec0 = DipolesSetBO.spec_sextupole
        xrms = self.currents_minmax()
        plt.plot(xrms, (+DipolesSetBO.spec_sextupole_rms,)*2, 'k--')
        plt.plot(xrms, (-DipolesSetBO.spec_sextupole_rms,)*2, 'k--')
        for i in range(len(magnets)):
            error = []
            for j in range(len(currents[i])):
                spec = spec0 * currents[i][j]/DipolesSetBO.spec_ref_current
                error.append(100 * (sextupoles[i][j] - spec)/spec)
            plt.plot(currents[i], error, 'o-', label=magnets[i])
        plt.xlabel('Current [A]')
        plt.ylabel('Sextupole Error w.r.t. to Spec [%]')
        plt.legend()

def get_summary(dst, cur, pos):
    """."""
    data = (
        'Booster Dipoles Integrated Principal Multipoles\n'
        '================================================\n'
        '\n'
        'As calculated in {0:s}-half Runge-Kutta trajectory,\n'
        'defined by measured fieldmap with magnet excitated with current of'
        ' {1:s},\n'
        'corresponding to nominal particle energy of 3 GeV.\n'
        '\n').format(pos, cur)

    fmt = '{0:^10s} | {1:^11s} |  {2:^11s} | {3:^11s} | {4:^11s} |\n'
    data += fmt.format('Dipole', 'Angle [°]', 'Dint [T.m]', 'Gint [T]',
                       'Sint [T/m]')
    data += fmt.format('', '', '', '', '')

    fmt = ('{0:^10s} | {1:^+11.5f} |  {2:^+11.5f} | {3:^+11.5f} |'
           ' {4:^+11.5f} |\n')
    for i in range(4, 58):
        name = 'bd-{0:03d}'.format(i)
        fi = dst + '/' + name + '/M1/{0:s}/z-{1:s}/'.format(cur, pos)
        fname = './' + fi + cur + '.pkl'
        # print(fname)
        with open(fname, 'rb') as fi:
            config = _pic.load(fi)
        si = 1 if pos == 'positive' else -1
        multi = config.multipoles.normal_multipoles_integral
        theta = si*_np.arctan(config.traj.px[-1]/config.traj.pz[-1])*180/_np.pi
        data += fmt.format(name, theta, multi[0], multi[1], multi[2])

    print(data)
    with open(dst + '/README-{0:s}-Z{1:s}.md'.format(cur, pos), 'w') as fi:
        fi.write(data)

def get_new_energy(dst, cur):
    """."""
    theta = _np.zeros(54)
    for i in range(4, 58):
        for pos in ['positive', 'negative']:
            fi = dst + '/bd-{0:03d}/M1/' + cur + '/z-{1:s}/'.format(i, pos)
            with open(fi + cur + '.pkl', 'rb') as fi:
                config = _pic.load(fi)
            si = 1 if pos == 'positive' else -1
            theta[i-4] += \
                si*_np.arctan(config.traj.px[-1]/config.traj.pz[-1])*180/_np.pi

    print(theta)
    print('New Energy = {0:7.5f} GeV'.format(theta.mean()/-7.2 * 3))

def print_help():
    """."""
    print(_help)

def main():
    """."""
    n = len(_sys.argv)
    if n > 1 and 'summary' in _sys.argv[1]:
        dst = _sys.argv[2] if n > 2 else 'x0-9p1013mm'
        cur = _sys.argv[3] if n > 3 else '0991p63A'
        pos = _sys.argv[4] if n > 4 else 'positive'
        get_summary(dst, cur, pos)
    elif n > 1 and 'energy' in _sys.argv[1]:
        dst = _sys.argv[2] if n > 2 else 'x0-9p1013mm'
        cur = _sys.argv[3] if n > 3 else '0991p63A'
        get_new_energy(dst, cur)
    elif n > 1 and 'help' in _sys.argv[1]:
        print_help()
    else:
        print_help()

if __name__ == '__main__':
    main()

max_stars_repo_name: carneirofc/lnls

max_issues_repo_name: carneirofc/lnls

max_forks_repo_name: carneirofc/lnls


text:
import numpy as np

def framework_candidate_select(framework_id, ref_dirs=None, pop=None, n_select=1):

    if n_select == 1:
        F = pop.get("F")
        if F.shape[1] > 1:  # if it is multi-objective, pick the middle one
            ref = (1 / F.shape[1]) * np.ones(F.shape[1])
            _F = np.sum(ref * F, 1)
            pop = pop[np.argmin(_F)]
        else:  # single-objective
            if np.any(pop.get("CV") <= 0):
                pop = pop[np.argmin(pop.get("F"))]
            else:
                pop = pop[np.argmin(pop.get("CV"))]
        return pop

    out_pop = []

    if framework_id in ['12', '22']:
        f = pop.get("F")
        f_asf = []
        if np.linalg.norm(np.max(f, axis=0) - np.min(f, axis=0)) > 1e-16:
            f_normalized = (f - np.min(f, axis=0)) / (np.max(f, axis=0) - np.min(f, axis=0))
        else:
            f_normalized = f

        for i in range(len(ref_dirs)):
            asf = np.max(f_normalized - ref_dirs[i], axis=1)
            f_asf.append(asf)

        f = np.column_stack(f_asf)
        cv = pop.get("CV")

    elif framework_id in ['32', '42']:
        f = pop.get("F")
        cv = pop.get("CV")
    else:
        f = pop.get("F")
        cv = np.zeros([f.shape[0], 1])

    feasible_index = (cv <= 0).flatten()  # find feasible index
    infeasible_index = (cv > 0).flatten()
    size = np.sum(infeasible_index)
    if size > 1:
        if np.sum(feasible_index) > 0:
            worst_of_all_ref_dir = np.max(f[feasible_index])
        else:
            worst_of_all_ref_dir = np.max(f)
        val_infeasible_sol = cv[infeasible_index] + worst_of_all_ref_dir
        f[infeasible_index, :] = np.tile(val_infeasible_sol, (1, f.shape[1]))

    for i in range(len(ref_dirs)):
        index = np.argsort(f[:, i])
        j = 0
        while index[j] in out_pop:
            j += 1
        out_pop.append(index[j])

    return pop[out_pop]

ext: py

max_stars_repo_name: proteekroy/samoo

max_issues_repo_name: proteekroy/samoo

max_forks_repo_name: proteekroy/samoo

reason: import numpy


text:
import numpy as np

from PDE._typing import *
from PDE.derivator import Derivator
from PDE import (
    utils,
    _internal
    )

class System:
    FDM_ORDER: set[int] = {2, 4}
    DERIVATIVE: set[int] = {1, 2}
    # (order, derivative) -> Coefs
    C_COEF: dict[tuple[int, int], list[float]] = {
        (2, 1): [-0.5, 0, 0.5],
        (2, 2): [1, 2, 1],
        (4, 1): [1/12, -2/3, 0, 2/3, -1/12],
        (4, 2): [-1/12, 4/3, -5/2, 4/3, -1/12]
    }
    A_COEF: dict[tuple[int, int], list[float]] = {
        (2, 1): [1.5, 2, 0.5],
        (2, 2): [-2, 5, -4, 1],
        (4, 1): [25/12, -4, 3, -4/3, 1/4],
        (4, 2): [-15/14, 77/6, -107/6, 13, -61/12]
    }

    def __init__(
            self,
            init_condition: FloatArray,
            boundary_condition: Optional[VecApply]=None,
            derive_order: int = 4,
            time_dif_order: int = 1
        ):
        pass

    @classmethod
    def create_derive(cls, order: int, derive_order: int, h: float,  space_shape: Sequence[int], dim: int=0) -> VecApply:
        if order not in cls.FDM_ORDER:
            raise ValueError(f"Order {order} not supported for derivation with FDM")

        invert_h = 1 / (h ** derive_order)
        c = np.array(cls.C_COEF[(order, derive_order)])
        a_f = np.array(cls.A_COEF[(order, derive_order)])
        a_b = -a_f[::-1]
        c_born = len(c) // 2
        a_born = len(a_b)

        def _derive(v: FloatArray) -> FloatArray:
            res = np.zeros(v.shape)
            # Centered
            for index, k in enumerate(range(-c_born, c_born + 1)):
                res = res + c[index] * np.roll(v, k, axis=dim)

            # forward/backward
            for f, b in utils.close_range(a_born):
                res[(slice(None),) * dim if dim else None, f, ...] = np.sum(a_f * np.take(v, np.arange(f, f + a_born), axis=dim))
                res[(slice(None),) * dim if dim else None, b, ...] = np.sum(a_b * np.take(v, np.arange(b - a_born + 1, b + 1), axis=dim))

            res = res * invert_h
            return res

        return _derive

max_stars_repo_name: zazbone/PDE

max_issues_repo_name: zazbone/PDE

max_forks_repo_name: zazbone/PDE


text:
import numpy as np
import torch
from .strategy import Strategy

class EntropySamplingDropout(Strategy):
    def __init__(self, dataset, net, n_drop=10):
        super(EntropySamplingDropout, self).__init__(dataset, net)
        self.n_drop = n_drop

    def query(self, n):
        unlabeled_idxs, unlabeled_data = self.dataset.get_unlabeled_data()
        probs = self.predict_prob_dropout(unlabeled_data, n_drop=self.n_drop)
        log_probs = torch.log(probs)
        uncertainties = (probs*log_probs).sum(1)
        return unlabeled_idxs[uncertainties.sort()[1][:n]]

max_stars_repo_name: nannullna/deep-active-learning

max_issues_repo_name: nannullna/deep-active-learning

max_forks_repo_name: nannullna/deep-active-learning


text:
import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = [14, 6]

st.title('Water Value Analysis')

# Filling levels
df = pd.read_csv('data/interim/fillinglevels.csv', index_col='acqdate')
df.index = pd.to_datetime(df.index)
df = df.resample('M').first()

# Model labels
mag_labels = np.array([1, 2, 0, 2, 0, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

# Cluster
st.subheader('Cluster')

st.write('Cluster 1: Blåsjø, Storglomvatn')
st.write('Cluster 2: Strandevatn, Roskreppfjord, Vatnedalsvatn, Aursjøen, Nesjø, Akersvatn, Altevatn, Sisovatn, Styggevatnet, Nyhellervatn, Svartevatn')
st.write('Cluster 3: Bygdin, Møsvatn, Sysenvatn, Tyin, Songavatn')

cluster_selection = st.selectbox(
    'Select cluster',
    (1, 2 , 3)
)

cluster = mag_labels == (cluster_selection - 1)
df = df[df.columns[cluster]]

# Water usage
st.subheader('Water Usage')
month_list = ['January', 'February', 'March', 'April', 'May', 'June', 'July', 'August', 'September', 'October', 'November', 'December']
start_month, end_month = st.select_slider(
    'Select range for water usage',
    options=month_list,
    value=('October', 'December'))
st.write('You have selected water usage between end of', start_month, 'and end of', end_month)

df_start = df[df.index.month == month_list.index(start_month)]
df_end = df[df.index.month == month_list.index(end_month)]

df_start = df_start.mean(axis=1)
df_end = df_end.mean(axis=1)

water_used = df_start.values - df_end.values

df_water_usage = pd.DataFrame(index=df_start.index.year)
df_water_usage['Water Used'] = water_used

# System price
st.subheader('System Price')
df_price = pd.read_csv('data/interim/system_price.csv', index_col='date')
df_price.index = pd.to_datetime(df_price.index)
df_price['year'] = df_price.index.year

week_list = list(range(1, 54))
start_week, end_week = st.select_slider(
    'Select system price range',
    options=week_list,
    value=(12, 16))
st.write('You have selected mean system price between week', start_week, 'and', end_week)

df_price = df_price[df_price.index.week.isin(list(range(start_week, end_week)))]
df_price = df_price.groupby('year').mean()

df_price.index = df_price.index + 1

# NVE
df_nve = pd.read_csv('data/interim/nve.csv', index_col='dato_Id')
df_nve.index = pd.to_datetime(df_nve.index)
df_nve = df_nve.resample('M').first()
df_nve = df_nve['fyllingsgrad']
df_nve_start = df_nve[df_nve.index.month == month_list.index(start_month)]
df_nve_end = df_nve[df_nve.index.month == month_list.index(end_month)]
df_nve_difference = df_nve_start.values - df_nve_end.values
df_nve = pd.DataFrame(index=df_nve_start.index.year)
df_nve['NVE'] = df_nve_difference

# Final DataFrame
df_water_usage['Price'] = df_price
final = df_water_usage.copy()
final = final.dropna()

# Plot
st.subheader('Result')

nve_choice = st.radio('Include NVE?',
('No', 'Yes'))

if nve_choice == 'Yes':
    final['NVE'] = df_nve
    final = final.dropna()

fig,ax = plt.subplots()
ax.set_title("Cluster {} \nWater used from {} to {} vs System Price between weeks {} and {}".format(cluster_selection, start_month, end_month, start_week, end_week), fontsize=14)
ax.plot((final['Water Used']), color="red", marker="o")
if nve_choice == 'Yes':
    ax.plot((final['NVE']), color="green", marker="o")
ax.set_ylabel("water used",color="red",fontsize=14)
ax2=ax.twinx()
ax2.plot(final['Price'], color="blue",marker="o")
ax2.set_ylabel("price [EUR/MWh]",color="blue",fontsize=14)
st.pyplot(fig)


text:
import numpy as np

# converts a complex 2d array into a 3d real array
def stackComplex(a):
	shape = (a.shape[0], a.shape[1], 1)
	re = np.reshape(a.real, shape)
	im = np.reshape(a.imag, shape)

	return np.concatenate((re,im), axis=2)

# inverse of stackComplex
def flattenComplex(a):
	if(a.shape[-1] == 1):
		a = np.reshape(a, a.shape[:-1])

	re = a[...,0]
	im = 1j * a[...,1]

	return re + im

def decompose(scalarField, lowPass):
	freq = np.fft.fft2(scalarField)
	# center
	freq = np.fft.fftshift(freq)
	begin = ((np.array(scalarField.shape) - lowPass) / 2).astype(int)
	end = (begin + lowPass).astype(int)
	lowFreq = freq[begin[0]:end[0],begin[1]:end[1]]

	return stackComplex(freq), stackComplex(lowFreq)

def invTransform(freqs):
	freqs = np.fft.ifftshift(freqs)

	return np.fft.ifft2(freqs)

def decomposeReal(scalarField, lowPass):
	freq = np.fft.rfftn(scalarField)
	freq = np.fft.fftshift(freq, axes=0)

	begin = (freq.shape[0] - lowPass[0]) // 2
	end = begin + lowPass[0]
	lowFreq = freq[begin:end,0:lowPass[1]]

	return stackComplex(freq), stackComplex(lowFreq)

def invTransformReal(freqs):
	flatFreqs = flattenComplex(freqs)
	flatFreqs = np.fft.ifftshift(flatFreqs,axes=0)
	return np.fft.irfftn(flatFreqs)

def shrink(freqs, lowPass):
	begin = (freqs.shape[0] - lowPass[0]) // 2
	end = begin + lowPass[0]

	redSize = freqs.shape[0]*freqs.shape[1] - lowPass[0] * lowPass[1]
	flatFreqs = np.zeros((redSize,2))
	ind = 0
	for y in range(freqs.shape[1]):
		for x in range(freqs.shape[0]):
			if x < begin or x >= end or y >= lowPass[1]:
				flatFreqs[ind,:] = freqs[x,y,:]
				ind += 1

	return flatFreqs

def composeReal(highFreqs, lowFreqs, highFreqShape):
	begin = (highFreqShape[0] - lowFreqs.shape[0]) // 2
	end = begin + lowFreqs.shape[0]

	freqs = np.zeros(highFreqShape)
	if(highFreqs.shape[-1] == 1):
		highFreqs = np.reshape(highFreqs, (highFreqs.shape[:-1]))

	ind = 0
	for y in range(highFreqShape[1]):
		for x in range(highFreqShape[0]):
			if x < begin or x >= end or y >= lowFreqs.shape[1]:
				freqs[x,y] = highFreqs[ind]
				ind += 1

	freqs[begin:end,0:lowFreqs.shape[1]] = lowFreqs

	return freqs

def getSpectrum(field, axis):
	field = flattenComplex(field)
	field = np.fft.ifftshift(field,axes=0)
	powerSpectrum = np.abs(field)**2

	if axis == 0:
		powerSpectrum = np.sum(powerSpectrum, axis=1)
		freqs = np.fft.fftfreq(field.shape[axis])
	elif axis == 1:
		powerSpectrum = np.sum(powerSpectrum, axis=0)
		freqs = np.fft.rfftfreq(field.shape[axis])

	idx = np.argsort(freqs)

	return freqs[idx], powerSpectrum[idx]

#	columns = np.fft.rfftfreq(freqs.shape[1])

max_stars_repo_name: Thanduriel/NeuralTurbulence

max_issues_repo_name: Thanduriel/NeuralTurbulence

max_forks_repo_name: Thanduriel/NeuralTurbulence


text:
from logging import debug

from numpy import append, array

from kuramoto_oscillator_model_perturbation import KuramotoOscillatorModelPerturbation

class KuramotoOscillatorModelNaturalFrequencyPerturbation(KuramotoOscillatorModelPerturbation):

    def __init__(self, start_time, affected_model, new_natural_frequency, end_time=None):
        KuramotoOscillatorModelPerturbation.__init__(self, start_time=start_time, affected_model=affected_model, end_time=end_time)

        old_natural_frequency = affected_model._get_current_setpoint_array()[0]

        if new_natural_frequency == old_natural_frequency:
            debug('New natural frequency matches the current one, change will be disabled')
            self.enabled = False
        else:
            self.old_natural_frequency = old_natural_frequency
            self.new_natural_frequency = new_natural_frequency

        self._change_type = 'kuramoto_oscillator_model_natural_frequency_perturbation'

    def _activate(self, t):
        # no need to try / except for the following method as only KuramotoOscillatorModel affected_models are allowed
        self.affected_model.change_dynamic_model_real_power_setpoint(self.new_natural_frequency)
        self.affected_model._setpoint_change_time = append(self.affected_model._setpoint_change_time, t)

    def _deactivate(self, t):
        self.affected_model.change_dynamic_model_real_power_setpoint(self.old_natural_frequency)
        self.affected_model._setpoint_change_time = append(self.affected_model._setpoint_change_time, t)

max_stars_repo_name: stantontcady/psyspy

max_issues_repo_name: stantontcady/psyspy

max_forks_repo_name: stantontcady/psyspy

reason: from numpy


text:
import tensorflow as tf
import numpy as np
import datetime
from .utils import echo

def simple_cov_matrix(
    blvecs,
    freqs,
    ant_dly=0.0,
    horizon=1.0,
    offset=0.0,
    min_dly=0.0,
    dtype=np.float64,
    use_tensorflow=False,
    verbose=False,
):
    """Compute simple covariance matrix for subset of baselines in uvdata

    Parameters
    ----------
    blvecs: list of np.ndarrays
        list of length-3 np.ndarrays representing ENH baseline vectors.
    freqs: array-like
        array of frequencies (Hz) sampled by interferometer.
    ant_dly: float
        intrinsic chromaticity of each antenna element, manifested in a
        multiplicative sinc matrix sinc(2 pi tau_ant * (nu_1 - nu_0))
    horizon: float, optional
        fraction of horizon for beam to extend to in analytic cov matrix.
        default is 1.0
    offset: float, optional
        additional offset causing additional decorrelation between antennas
        (units of ns).
        default is 0.0
    min_dly: float, optional
        minimum decorrelation delay between points in uv plane. Net effect should
        be some additional modes and signal loss.
        default is 0.0
    dtype: numpy.dtype
        data type to process and store covariance matrix as.
    use_tensorflow: bool, optional
        if True, use tensorflow for multi-baseline modeling matrix operations.
        Recommended only for machines with GPUs.
    verbose: bool, optional
        lots of text outputs.
        default is False.

    Returns
    -------
    cmat: tf.Tensor object (if use_tensorflow is True) or np.ndarray(use_tensorflow is False)
        (Nbls * Nfreqs) x (Nbls * Nfreqs) tf.Tensor object or np.ndarray with dtype dtype

    """
    nbls = len(blvecs)
    uvws = np.asarray(blvecs, dtype=dtype)
    freqs = np.asarray(freqs, dtype=dtype)
    nfreqs = len(freqs)
    freqs = freqs.astype(dtype)
    absdiff = np.zeros((nbls * nfreqs, nbls * nfreqs), dtype=dtype)
    if use_tensorflow:
        freqs = tf.convert_to_tensor(freqs, dtype=dtype)
        absdiff = tf.convert_to_tensor(absdiff, dtype=dtype)
    for uvw_ind in range(3):
        if use_tensorflow:
            uvw_coord = tf.reshape(tf.experimental.numpy.outer(uvws[:, uvw_ind], freqs / 3e8), (nbls * nfreqs,))
            cg0, cg1 = tf.meshgrid(uvw_coord, uvw_coord, indexing="ij")
            absdiff += tf.math.square(tf.math.abs(cg0 - cg1) * horizon)
        else:
            uvw_coord = np.reshape(np.outer(uvws[:, uvw_ind], freqs / 3e8), (nbls * nfreqs,))
            cg0, cg1 = np.meshgrid(uvw_coord, uvw_coord, indexing="ij")
            absdiff += np.abs(cg0 - cg1) ** 2.0
    if use_tensorflow:
        absdiff = tf.math.sqrt(absdiff)
        fvals = tf.reshape(tf.experimental.numpy.outer(tf.ones(nbls, dtype=dtype), freqs), (nbls * nfreqs,))
        fg0, fg1 = tf.meshgrid(fvals, fvals, indexing="ij")
        dfg = tf.abs(fg0 - fg1) / 1e9
    else:
        absdiff = np.sqrt(absdiff) * horizon
        fvals = np.reshape(np.outer(np.ones(nbls, dtype=dtype), freqs), (nbls * nfreqs,))
        fg0, fg1 = np.meshgrid(fvals, fvals, indexing="ij")
        dfg = np.abs(fg0 - fg1) / 1e9

    del fg0, fg1
    absdiff += dfg * offset
    if use_tensorflow:
        cmat = tf.experimental.numpy.sinc(2 * tf.maximum(absdiff, min_dly * dfg))
    else:
        cmat = np.sinc(2 * np.maximum(min_dly * dfg, absdiff))
    del absdiff
    if use_tensorflow:
        cmat = cmat * tf.experimental.numpy.sinc(2 * dfg * ant_dly)
    else:
        cmat = cmat * np.sinc(2 * dfg * ant_dly)
    del dfg
    return cmat

def yield_simple_multi_baseline_model_comps(
    blvecs,
    freqs,
    ant_dly=0.0,
    horizon=1.0,
    offset=0.0,
    min_dly=0.0,
    dtype=np.float64,
    verbose=False,
    use_tensorflow=False,
    eigenval_cutoff=1e-10,
):
    """Generate model components for multiple baselines.

    Parameters
    ----------
    blvecs: list of len-3 numpy.ndarrays
      list of baseline vectors ENH
    ant_dly: float, optional
        intrinsic chromaticity of each antenna element, manifested in a
        multiplicative sinc matrix sinc(2 pi tau_ant * (nu_1 - nu_0))
        default is 0.
    freqs: array-like
      array of frequencies sampled by interferometer (MHz).
    horizon: float, optional
        fraction of horizon for beam to extend to in analytic cov matrix.
        default is 1.0
    offset: float, optional
        additional offset causing additional decorrelation between antennas
        (units of ns).
        default is 0.0
    min_dly: float, optional
        minimum decorrelation delay between points in uv plane. Net effect should
        be some additional modes and signal loss.
        default is 0.0
    dtype: numpy.dtype, optional
        data-type to use for eigenvalue decomposition
        default is np.float32
    verbose: bool, optional
        text outputs
        default is False.
    use_tensorflow: bool, optional
        if True, use tensorflow for multi-baseline modeling matrix operations.
        Recommended only for machines with GPUs.
    eigenval_cutoff: float, optional
        threshold of eigenvectors to include in modeling components.

    Returns
    -------
    multi_baseline_model_comps: np.ndarray
        (Nbl * Nfreqs) x Ncomponents numpy.ndarray containing eigenvectors
        with eigenvalues above cutoff threshold.

    """
    cmat = simple_cov_matrix(
        blvecs,
        freqs,
        ant_dly=ant_dly,
        horizon=horizon,
        offset=offset,
        min_dly=min_dly,
        dtype=dtype,
        verbose=verbose,
        use_tensorflow=use_tensorflow,
    )
    # get eigenvectors
    echo(
        f"{datetime.datetime.now()} Deriving modeling components with eigenvalue decomposition...\n",
        verbose=verbose,
    )
    if use_tensorflow:
        evals, evecs = tf.linalg.eigh(cmat)
        evals = evals.numpy()
        evecs = evecs.numpy()
    else:
        evals, evecs = np.linalg.eigh(cmat)
    selection = evals / evals[-1] >= eigenval_cutoff
    echo(
        f"{datetime.datetime.now()} Finished spectral decomposition. Using {np.count_nonzero(selection)} of {len(selection)} eigenvectors to model foregrounds...\n"
    )
    evals = evals[selection][::-1]
    evecs = evecs[:, selection][:, ::-1]
    return evecs

max_stars_repo_name: LBJ-Wade/calamity

max_issues_repo_name: LBJ-Wade/calamity

max_forks_repo_name: LBJ-Wade/calamity


text:
"""Cuda op Python library."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os.path
import numpy as np
import tensorflow.compat.v1 as tf
from tensorflow.python.framework import ops
from tensorflow.python.training  import slot_creator
from tensorflow.python.training  import optimizer
from blocksparse.utils import _op_module, scalar_constant, is_param_casted
from blocksparse.ewops import float_cast
from blocksparse.quantize import quantize

############################## AdamOptimizer #####################################

adam_op = _op_module.adam
blocksparse_adam_op =_op_module.blocksparse_adam

class AdamOptimizer(optimizer.Optimizer):
    def __init__(self,
        learning_rate=3e-4, beta1=0.9, beta2=0.999, epsilon=1e-8, clip_sigmas=0.0,
        norm_scale=None, grad_scale=1.0, saturate=0.0, zero_infs=False, zero_nans=False,
        gated=False, param_qspec=None, mean_qspec=None, var_qspec=None,
        fp16=False, zero_init_variables=False, name="Adam"):

        super().__init__(False, name)
        self.beta1       = beta1
        self.beta2       = beta2
        self.epsilon     = epsilon
        self.saturate    = saturate
        self.zero_infs   = zero_infs
        self.zero_nans   = zero_nans
        self.gated       = gated
        self.param_qspec = param_qspec
        self.mean_qspec  = mean_qspec
        self.var_qspec   = var_qspec
        self.name        = name
        self.norm_scale  = [] if norm_scale is None else [norm_scale]
        self.fp16        = fp16

        beta1_init = 0.0 if zero_init_variables else beta1
        beta2_init = 0.0 if zero_init_variables else beta2

        with tf.device("/cpu:0"), tf.variable_scope("adam_beta"):

            one = scalar_constant(1.0, dtype=tf.float32)
            self.beta1_power = tf.Variable(initial_value=beta1_init, name="beta1_power", trainable=False)
            self.beta2_power = tf.Variable(initial_value=beta2_init, name="beta2_power", trainable=False)
            self.beta1_t     = scalar_constant(beta1,         dtype=tf.float32)
            self.beta2_t     = scalar_constant(beta2,         dtype=tf.float32)
            self.clip_sigma  = scalar_constant(clip_sigmas,   dtype=tf.float32)
            self.grad_scale  = scalar_constant(grad_scale,    dtype=tf.float32)
            self.lr          = scalar_constant(learning_rate, dtype=tf.float32) * tf.sqrt(one - self.beta2_power) / (one - self.beta1_power)

    def _get_beta_accumulators(self):
        return self.beta1_power, self.beta2_power

    def _non_slot_variables(self):
        return self._get_beta_accumulators()

    def _create_slots(self, params):
        # Create slots for the first and second moments.
        with tf.device("/gpu:0"), tf.control_dependencies(None):
            for param in params:
                # only use fp16 for larger params that benefit from memory savings
                dtype = tf.float16 if self.fp16 and param.shape.num_elements() >= 8*1024 else tf.float32  #is_param_casted(param)

                self._get_or_make_slot(param, tf.zeros(param.shape, dtype=dtype), "Mean", self.name)
                self._get_or_make_slot(param, tf.zeros(param.shape, dtype=dtype), "Var", self.name)

    def _apply_dense(self, grad, param):

        m = self.get_slot(param, "Mean")
        v = self.get_slot(param, "Var")

        gate = getattr(param, "gate", None)
        gate = [gate] if self.gated and gate is not None else []

        op = adam_op(grad, param, m, v, self.lr, self.grad_scale, self.clip_sigma, self.norm_scale, gate,
            decay_mean=self.beta1, decay_var=self.beta2, epsilon=self.epsilon,
            saturate=self.saturate, zero_infs=self.zero_infs, zero_nans=self.zero_nans, lazy_emb=hasattr(grad, "lazy"))

        updates = list()
        if self.param_qspec is not None:
            updates.append(param.assign(quantize(op.out_param, self.param_qspec, name="param_" + param.op.name)))
        else:
            updates.append(op.out_param)

        if self.mean_qspec is not None:
            updates.append(m.assign(quantize(op.out_mean, self.mean_qspec, name="mean_" + param.op.name)))

        if self.var_qspec is not None:
            updates.append(v.assign(quantize(op.out_var, self.var_qspec, name="var_" + param.op.name)))

        return tf.group(*updates) if len(updates) > 1 else updates[0]

    def _apply_sparse(self, grad, param):
        raise NotImplementedError("Sparse gradient updates are not supported.")

    def _finish(self, update_ops, name_scope):
        # Update the power accumulators.
        with ops.control_dependencies([ self.lr ]), tf.device("/cpu:0"):
            update_beta1 = self.beta1_power.assign(self.beta1_power * self.beta1_t)
            update_beta2 = self.beta2_power.assign(self.beta2_power * self.beta2_t)

        return tf.group(*update_ops + [update_beta1, update_beta2], name=name_scope)

############################## AdafactorOptimizer #####################################

adafactor1d_op = _op_module.adafactor1d
adafactor2d_op = _op_module.adafactor2d

class AdafactorOptimizer(optimizer.Optimizer):
    def __init__(self, learning_rate=5e-4, beta2=0.999, epsilon=1e-30, clip_thresh=1.0,
        norm_scale=None, grad_scale=1.0, saturate=0.0, zero_infs=False, zero_nans=False,
        name="Adafactor", zero_init_variables=False):

        super().__init__(False, name)
        self.epsilon    = epsilon
        self.saturate   = saturate
        self.zero_infs  = zero_infs
        self.zero_nans  = zero_nans
        self.name       = name
        self.norm_scale = [] if norm_scale is None else [norm_scale]

        beta2_init = 0.0 if zero_init_variables else beta2

        with tf.device("/cpu:0"), tf.variable_scope("adafactor_decay"):

            one = scalar_constant(1.0, dtype=tf.float32)
            self.decay1_power = tf.Variable(initial_value=beta2_init,            name="decay1_power", trainable=False)
            self.decay2_power = tf.Variable(initial_value=beta2_init*beta2_init, name="decay2_power", trainable=False)
            self.learn_rate   = scalar_constant(learning_rate, dtype=tf.float32)
            self.clip_thresh  = scalar_constant(clip_thresh,   dtype=tf.float32)
            self.grad_scale   = scalar_constant(grad_scale,    dtype=tf.float32)
            self.decay_t      = scalar_constant(beta2,         dtype=tf.float32)
            self.decay        = self.decay_t * (one - self.decay1_power) / (one - self.decay2_power)

    def _get_beta_accumulators(self):
        return self.decay1_power, self.decay2_power

    def _non_slot_variables(self):
        return self._get_beta_accumulators()

    def _create_slots(self, params):
        # Create slots for the first and second moments.
        for param in params:
            if param.shape.ndims == 2 and param.shape[0].value > 1:
                self._get_or_make_slot(param, tf.zeros(param.shape[1].value), "cv", self.name + "CV")
                self._get_or_make_slot(param, tf.zeros(param.shape[0].value), "rv", self.name + "RV")
            elif param.shape.ndims == 1 or (param.shape.ndims == 2 and param.shape[0].value == 1):
                self._get_or_make_slot(param, tf.zeros(param.shape.num_elements()), "cv", self.name + "CV")
            else:
                raise ValueError("only 1 or 2d params are supported")

    def _apply_dense(self, grad, param):

        if param.shape.ndims == 2 and param.shape[0].value > 1:

            cv = self.get_slot(param, "cv")
            rv = self.get_slot(param, "rv")

            return adafactor2d_op(param, cv, rv, grad,
                self.decay, self.learn_rate, self.grad_scale, self.clip_thresh, self.norm_scale, epsilon=self.epsilon,
                saturate=self.saturate, zero_infs=self.zero_infs, zero_nans=self.zero_nans).out_param

        elif param.shape.ndims == 1 or (param.shape.ndims == 2 and param.shape[0].value == 1):

            cv = self.get_slot(param, "cv")

            return adafactor1d_op(param, cv, grad,
                self.decay, self.learn_rate, self.grad_scale, self.clip_thresh, self.norm_scale, epsilon=self.epsilon,
                saturate=self.saturate, zero_infs=self.zero_infs, zero_nans=self.zero_nans).out_param
        else:
            raise ValueError("only 1 or 2d params are supported")

    def _apply_sparse(self, grad, param):
        raise NotImplementedError("Sparse gradient updates are not supported.")

    def _finish(self, update_ops, name_scope):
        # Update the power accumulators.
        with ops.control_dependencies([ self.decay ]), tf.device("/cpu:0"):
            update_decay1 = self.decay1_power.assign(self.decay1_power * self.decay_t)
            update_decay2 = self.decay2_power.assign(self.decay2_power * self.decay_t)

        return tf.group(*update_ops + [update_decay1, update_decay2], name=name_scope)

############################## ClipGlobalNorm #####################################

clip_global_norm_op = _op_module.clip_global_norm

def clip_by_global_norm(grads, clip_norm=1.0, grad_scale=1.0, saturate=0.0, zero_infs=False, zero_nans=False):

    grad_float = list()
    grad_ehalf = list()
    grad_bhalf = list()

    for grad in grads:
        if   grad.dtype is tf.float32:
            grad_float.append(grad)
        elif grad.dtype is tf.float16:
            grad_ehalf.append(grad)
        elif grad.dtype is tf.bfloat16:
            grad_bhalf.append(grad)
        else:
            raise ValueError("unsupported grad dtype")

    with tf.device("/gpu:0"):
        global_norm, norm_scale, _ = clip_global_norm_op(
            scalar_constant(grad_scale, dtype=tf.float32),
            scalar_constant(clip_norm,  dtype=tf.float32),
            grad_float, grad_ehalf, grad_bhalf,
            saturate=saturate, zero_infs=zero_infs, zero_nans=zero_nans)

    return global_norm, norm_scale

def global_norm(grads, grad_scale=1.0, saturate=0.0, zero_infs=False, zero_nans=False):
    gn, _ = clip_by_global_norm(grads, clip_norm=9e9, grad_scale=grad_scale, saturate=saturate, zero_infs=zero_infs, zero_nans=zero_nans)
    return gn

# old function name
def ClipGlobalNorm(grads, clip_norm=1.0, grad_scale=1.0, saturate=0.0, zero_infs=False, zero_nans=False):
    return clip_by_global_norm(grads, clip_norm=clip_norm, grad_scale=grad_scale, saturate=saturate, zero_infs=zero_infs, zero_nans=zero_nans)

############################## Exponential Moving Average #####################################

ema_op = _op_module.ema

class Ema(object):

    def __init__(self, decay=0.999, gated=False, fp16=False, name="Ema"):
        self.decay    = decay
        self.gated    = gated
        self.fp16     = fp16
        self.name     = name
        self.averages = dict()

    def apply(self, params, qspec=None):

        with tf.device("/gpu:0"), tf.control_dependencies(None):
            for param in params:
                if self.fp16 == 2 or (self.fp16 and is_param_casted(param)):
                    # only use fp16 for params that are explicitly cast to fp16 before use
                    init  = float_cast(param.initialized_value(), dtype=tf.float16)
                    dtype = tf.float16
                else:
                    init  = param.initialized_value()
                    dtype = tf.float32

                with tf.variable_scope(None, param.op.name + "/" + self.name):
                    # use the Identity read op output as the key
                    # this lets us lookup ema vars by Cast op outputs
                    self.averages[param.value()] = tf.get_variable("ema", dtype=dtype, initializer=init, trainable=False)
                ops.add_to_collection(ops.GraphKeys.MOVING_AVERAGE_VARIABLES, param)

        ema_ops = []
        for param in params:

            ema   = self.averages[param.value()]
            gate  = getattr(param, "gate", None)
            gate  = [gate] if self.gated and gate is not None else []

            op = ema_op(ema, param, gate, decay=self.decay)

            if qspec is not None:
                ema_ops.append(ema.assign(quantize(op, qspec, name="ema_" + param.op.name)))
            else:
                ema_ops.append(op)

        return tf.group(*ema_ops)

    def average(self, param):
        if isinstance(param, tf.Variable):
            # this is just a raw param
            key = param.value()
        elif isinstance(param, tf.Tensor):
            # we're given a Cast op output
            # TODO: maybe traverse deeper?
            key = param.op.inputs[0]
        else:
            raise TypeError("bad param type")

        return self.averages.get(key, None)

############################## Group LASSO / Blocksparse L2 decay #####################################

l2_decay_op = _op_module.blocksparse_l2_decay

def _check_param_shape(param, gate=None):
    assert len(param.shape) == 3 and param.shape[1].value == param.shape[2].value and param.shape[1].value in (8,16,32,64)
    if gate is not None:
        assert gate.shape.num_elements() == param.shape[0].value

def blocksparse_l2_decay(param, gate=None, rate=0.05, epsilon=1e-12):

    _check_param_shape(param, gate)

    gate = [gate] if gate is not None else []

    return l2_decay_op(param, scalar_constant(rate, dtype=tf.float32), gate, epsilon=epsilon)

############################## Blocksparse Pruning #####################################

blocksparse_norm_op            = _op_module.blocksparse_norm
blocksparse_prune_op           = _op_module.blocksparse_prune
blocksparse_threshold_prune_op = _op_module.blocksparse_threshold_prune

def blocksparse_norm(param, norm="max"):
    _check_param_shape(param)
    return blocksparse_norm_op(param, norm_type=1 if norm.lower() == "l2" else 0)

def blocksparse_prune(param, gate, step, sparsity=None, threshold=None, norm="max", frequency=1):

    _check_param_shape(param, gate)

    # one must be set
    assert (sparsity is None) ^ (threshold is None)

    if sparsity is not None:

        # apply pruning to the moving average
        norms = blocksparse_norm(param, norm=norm)

        k = scalar_constant(param.shape[0].value, dtype=tf.int32)

        _, idx = tf.nn.top_k(norms, k=k, sorted=True)

        return blocksparse_prune_op(gate, idx, scalar_constant(sparsity, dtype=tf.float32), step, frequency=frequency)

    elif threshold is not None:

        norm = 1 if norm.lower() == "l2" else 0

        return blocksparse_threshold_prune_op(gate, param, scalar_constant(threshold, dtype=tf.float32), step, frequency=frequency, norm_type=norm)

max_stars_repo_name: AlumiK/blocksparse

max_issues_repo_name: AlumiK/blocksparse

max_forks_repo_name: AlumiK/blocksparse


text:
import numpy as np

from tengp_eval.utils.symreg import get_benchmark_keijzer

def get_data(generator, random, bench_id):
    """
    get data from given generator and add constant column to inputs
    """
    x_train, y_train, x_test, y_test = generator(random, bench_id)
    x_train = np.c_[np.ones(len(x_train)), x_train]
    x_test = np.c_[np.ones(len(x_test)), x_test]
    return x_train, y_train, x_test, y_test

def get_keijzer_data(random, bench_id):
    return get_data(get_benchmark_keijzer, random, bench_id)

max_stars_repo_name: Jarino/cgp-optimization

max_issues_repo_name: Jarino/cgp-optimization

max_forks_repo_name: Jarino/cgp-optimization


text:
import numpy as np
import pymc3 as pm
import theano.tensor as tt

from sampled import sampled
SEED = 1

def test_sampled_one_model():
    @sampled
    def just_a_normal():
        pm.Normal('x', mu=0, sd=1)

    kwargs = {
        'draws': 50,
        'tune': 50,
        'init': None
    }

    np.random.seed(SEED)
    with just_a_normal():
        decorated_trace = pm.sample(**kwargs)

    np.random.seed(SEED)
    with pm.Model():
        pm.Normal('x', mu=0, sd=1)
        normal_trace = pm.sample(**kwargs)

    assert decorated_trace.varnames == ['x']
    np.testing.assert_array_equal(decorated_trace.get_values('x'), normal_trace.get_values('x'))

def test_reuse_model():
    @sampled
    def two_normals():
        mu = pm.Normal('mu', mu=0, sd=1)
        pm.Normal('x', mu=mu, sd=1)

    with two_normals():
        generated_data = pm.sample(draws=50, tune=50, init=None)

    for varname in ('mu', 'x'):
        assert varname in generated_data.varnames

    with two_normals(mu=1):
        posterior_data = pm.sample(draws=50, tune=50, init=None)

    assert 'x' in posterior_data.varnames
    assert 'mu' not in posterior_data.varnames
    assert posterior_data.get_values('x').mean() > generated_data.get_values('x').mean()

def test_linear_model():
    rows, cols = 100, 10
    X = np.random.normal(size=(rows, cols))
    w = np.random.normal(size=cols)
    y = X.dot(w) + np.random.normal(scale=0.1, size=rows)

    @sampled
    def linear_model(X, y):
        shape = X.shape
        X = pm.Normal('X', mu=np.mean(X, axis=0), sd=np.std(X, axis=0), shape=shape)
        coefs = pm.Normal('coefs', mu=tt.zeros(shape[1]), sd=tt.ones(shape[1]), shape=shape[1])
        pm.Normal('y', mu=tt.dot(X, coefs), sd=tt.ones(shape[0]), shape=shape[0])

    with linear_model(X=X, y=y):
        sampled_coefs = pm.sample(draws=1000, tune=500)
    mean_coefs = sampled_coefs.get_values('coefs').mean(axis=0)
    np.testing.assert_allclose(mean_coefs, w, atol=0.1)

def test_partial_model():
    rows, cols = 100, 10
    X = np.random.normal(size=(rows, cols))
    w = np.random.normal(size=cols)
    y = X.dot(w) + np.random.normal(scale=0.1, size=rows)

    @sampled
    def partial_linear_model(X):
        shape = X.shape
        X = pm.Normal('X', mu=np.mean(X, axis=0), sd=np.std(X, axis=0), shape=shape)
        pm.Normal('coefs', mu=tt.zeros(shape[1]), sd=tt.ones(shape[1]), shape=shape[1])

    with partial_linear_model(X=X) as model:
        coefs = model.named_vars['coefs']
        pm.Normal('y', mu=tt.dot(X, coefs), sd=tt.ones(y.shape), observed=y)
        sampled_coefs = pm.sample(draws=1000, tune=500)

    mean_coefs = sampled_coefs.get_values('coefs').mean(axis=0)
    np.testing.assert_allclose(mean_coefs, w, atol=0.1)

max_stars_repo_name: ColCarroll/sampled

max_issues_repo_name: ColCarroll/sampled

max_forks_repo_name: ColCarroll/sampled

reason: import numpy,import theano,import pymc3


text:
"""A script to compare the feature data collected."""

from collections import defaultdict, Counter
from functools import partial

import numpy as np

from classes import classes
from models import Repo

def summarize_imports(class_map):
    counter_map = defaultdict(Counter)
    for clsname, repos in class_map.items():
        for r in repos:
            counter_map[clsname].update(r.imported_stdlib_modules)

    # normalize to percent occurance
    for clsname, counter in counter_map.items():
        num_repos = float(len(class_map[clsname]))
        for mod in counter:
            counter[mod] /= num_repos

    print 'top imports'
    for clsname, counter in counter_map.items():
        print clsname
        for mod, occurs in counter.most_common(20):
            print "    {} {:.2g}".format(mod, occurs)
        print
    print

    if len(class_map) == 2:
        print 'largest absolute differences'
        a_name, b_name = counter_map.keys()
        print "(left is %s)" % a_name
        print

        differences = Counter()
        differences.update(counter_map[a_name])
        differences.subtract(counter_map[b_name])

        for mod in differences:
            differences[mod] = abs(differences[mod])

        for mod, delta in differences.most_common(20):
            print " {} {:.2g}".format(mod, delta)
            print "    {:.2g} {:.2g}".format(
                counter_map[a_name][mod], counter_map[b_name][mod])

    print

    if len(class_map) == 2:
        print 'relative differences'
        print '(only modules in both classes are included)'
        a_name, b_name = counter_map.keys()
        a_cnt, b_cnt = counter_map[a_name], counter_map[b_name]
        print "(left is %s)" % a_name
        print

        deltas = []  # (import, factor, more_frequent_class)
        for mod in (counter_map[a_name].viewkeys() &
                    counter_map[b_name].viewkeys()):
            # ugh
            vals = zip((a_cnt[mod], b_cnt[mod]), (a_name, b_name))
            big_val, big_class_name = max(vals, key=lambda p: p[0])
            small_val, small_class_name = min(vals, key=lambda p: p[0])

            deltas.append((mod, 1.0 * small_val / big_val, big_class_name))

        deltas.sort(key=lambda t: t[1])

        def output(deltas):
            for mod, rel_delta, big_class_name in deltas:
                print " {} {:.2g}".format(mod, rel_delta)
                print "    {:.2g} {:.2g}".format(
                    counter_map[a_name][mod], counter_map[b_name][mod])

        print 'smallest'
        output(deltas[-20:])
        print

        print 'largest'
        output(deltas[:20])

percentiles = [partial(np.percentile, q=i) for i in (1, 25, 50, 75, 99)]
default_funcs = [np.amin] + percentiles + [np.amax, np.mean, np.std]

def summarize_features(class_map, feature_names, class_names, funcs=None):
    if funcs is None:
        funcs = default_funcs

    feature_summaries = {}  # eg 'ReadmeSize' -> np.array(num_classes, num_funcs)

    for feature_name in feature_names:
        class_summaries = []
        for class_name in class_names:
            repos = class_map[class_name]
            feature_data = [getattr(repo, feature_name) for repo in repos]

            class_summaries.append([f(feature_data) for f in funcs])

        feature_summaries[feature_name] = np.array(class_summaries)

    for feature_name, summary in feature_summaries.items():
        print feature_name
        for i, class_name in enumerate(class_names):
            print '  ', class_name
            print '    ', '\t'.join(str(e) for e in summary[i])

        print '-----'
        print

if __name__ == '__main__':
    class_map = Repo.load_sample(separate=True)

    summarize_features(class_map, ['with_stmt_usage',
                                   'compr_usage',
                                   'lambda_usage',
                                   'global_usage',
                                   'gen_exp_usage',
                                   'print_usage',
                                  ], sorted(classes))

    #summarize_imports(class_map)

max_stars_repo_name: simon-weber/Predicting-Code-Popularity

max_issues_repo_name: simon-weber/Predicting-Code-Popularity

max_forks_repo_name: simon-weber/Predicting-Code-Popularity


text:
# -*- coding: utf-8 -*-
"""
Created on Wed Nov 29 09:52:35 2017

@author: gregz
"""
import glob
import re

import tables as tb
import argparse as ap
import os.path as op
import numpy as np

from astropy.io import fits
from input_utils import setup_logging
from astropy.table import Table

def build_path(reduction_folder, instr, date, obsid, expn):
    folder = op.join(date, instr, "{:s}{:07d}".format(instr, int(obsid)),
                     "exp{:02d}".format(int(expn)), instr)
    return op.join(reduction_folder, folder)

def get_files(args):
    files = glob.glob(op.join(args.rootdir, args.date, 'virus',
                              'virus%07d' % int(args.observation),
                              'exp*', 'virus', 'multi_*.fits'))
    return files

class VIRUSFiber(tb.IsDescription):
    obsind = tb.Int32Col()
    fibnum = tb.Int32Col()
    ifux = tb.Float32Col()
    ifuy = tb.Float32Col()
    fpx = tb.Float32Col()
    fpy = tb.Float32Col()
    ra = tb.Float32Col()
    dec = tb.Float32Col()
    spectrum = tb.Float32Col((1032,))
    wavelength = tb.Float32Col((1032,))
    fiber_to_fiber = tb.Float32Col((1032,))
    twi_spectrum = tb.Float32Col((1032,))
    trace = tb.Float32Col((1032,))
    sky_subtracted = tb.Float32Col((1032,))
    error1Dfib = tb.Float32Col((1032,))
    ifuslot = tb.StringCol(3)
    ifuid = tb.StringCol(3)
    specid = tb.StringCol(3)
    amp = tb.StringCol(2)
    expnum = tb.Int32Col()

class VIRUSImage(tb.IsDescription):
    obsind = tb.Int32Col()
    image = tb.Float32Col((1032, 1032))
    error = tb.Float32Col((1032, 1032))
    clean_image = tb.Float32Col((1032, 1032))
    ifuslot = tb.StringCol(3)
    ifuid = tb.StringCol(3)
    specid = tb.StringCol(3)
    amp = tb.StringCol(2)
    expnum = tb.Int32Col()

class VIRUSShot(tb.IsDescription):
    obsind = tb.Int32Col()
    date = tb.StringCol(8)
    mjd = tb.Float32Col()
    obsid = tb.StringCol(7)
    ra = tb.Float32Col()
    dec = tb.Float32Col()
    pa = tb.Float32Col()
    expn = tb.Int32Col()
    time = tb.StringCol(7)
    ambtemp = tb.Float32Col()
    humidity = tb.Float32Col()
    dewpoint = tb.Float32Col()
    pressure = tb.Float32Col()
    exptime = tb.Float32Col()

def append_shot_to_table(shot, fn, cnt):
    F = fits.open(fn)
    shot['obsind'] = cnt
    shot['date'] = ''.join(F[0].header['DATE-OBS'].split('-'))
    shot['time'] = ''.join(re.split('[:,.]', F[0].header['UT']))[:7]
    shot['mjd'] = F[0].header['MJD']
    shot['obsid'] = '%07d' % F[0].header['OBSID']
    shot['ra'] = F[0].header['TRAJCRA'] * 15.
    shot['dec'] = F[0].header['TRAJCDEC']
    shot['pa'] = F[0].header['PARANGLE']
    shot['ambtemp'] = F[0].header['AMBTEMP']
    shot['humidity'] = F[0].header['HUMIDITY']
    shot['dewpoint'] = F[0].header['DEWPOINT']
    shot['pressure'] = F[0].header['BAROMPRE']
    shot['exptime'] = F[0].header['EXPTIME']
    shot['expn'] = int(op.basename(op.dirname(op.dirname(F.filename())))[-2:])
    shot.append()

def append_fibers_to_table(fib, im, fn, cnt, T):
    F = fits.open(fn)
    n = F['spectrum'].data.shape[0]
    d = F['spectrum'].data.shape[1]
    attr = ['spectrum', 'wavelength', 'fiber_to_fiber', 'twi_spectrum',
            'sky_subtracted', 'trace', 'error1Dfib']
    imattr = ['image', 'error', 'clean_image']
    for att in imattr:
        if att == 'image':
            im[att] = F['PRIMARY'].data * 1.
        else:
            im[att] = F[att].data * 1.
    mname = op.basename(fn)[:-5]
    expn = op.basename(op.dirname(op.dirname(fn)))
    sel = T['col8'] == (mname + '_001.ixy')
    sel1 = T['col10'] == expn
    loc = np.where(sel * sel1)[0]
    for i in np.arange(n):
        fib['obsind'] = cnt
        fib['fibnum'] = i
        loci = loc + i
        if len(loc):
            fib['ra'] = T['col1'][loci]
            fib['dec'] = T['col2'][loci]
            fib['fpx'] = T['col6'][loci]
            fib['fpy'] = T['col7'][loci]
        else:
            fib['ra'] = -999.0
            fib['dec'] = -999.0
            fib['fpx'] = -999.0
            fib['fpy'] = -999.0
        fib['ifux'] = F['ifupos'].data[i, 0]
        fib['ifuy'] = F['ifupos'].data[i, 1]
        for att in attr:
            if att in F:
                fib[att] = F[att].data[i, :]
            else:
                fib[att] = np.zeros((d,))
        fib['ifuslot'] = '%03d' % int(F[0].header['IFUSLOT'])
        fib['ifuid'] = '%03d' % int(F[0].header['IFUID'])
        fib['specid'] = '%03d' % int(F[0].header['SPECID'])
        fib['amp'] = '%s' % F[0].header['amp'][:2]
        fib['expnum'] = int(expn[-2:])
        fib.append()
    im['obsind'] = cnt
    im['ifuslot'] = '%03d' % int(F[0].header['IFUSLOT'])
    im['ifuid'] = '%03d' % int(F[0].header['IFUID'])
    im['specid'] = '%03d' % int(F[0].header['SPECID'])
    im['amp'] = '%s' % F[0].header['amp'][:2]
    im['expnum'] = int(expn[-2:])
    im.append()
    return True

def main(argv=None):
    ''' Main Function '''
    # Call initial parser from init_utils
    parser = ap.ArgumentParser(description="""Create HDF5 file.""",
                               add_help=True)

    parser.add_argument("-d", "--date",
                        help='''Date, e.g., 20170321, YYYYMMDD''',
                        type=str, default=None)

    parser.add_argument("-o", "--observation",
                        help='''Observation number, "00000007" or "7"''',
                        type=str, default=None)

    parser.add_argument("-r", "--rootdir",
                        help='''Root Directory for Reductions''',
                        type=str, default='/work/03946/hetdex/maverick')

    parser.add_argument('-of', '--outfilename', type=str,
                        help='''Relative or absolute path for output HDF5
                        file.''', default=None)
    parser.add_argument('-a', '--append',
                        help='''Appending to existing file.''',
                        action="count", default=0)

    args = parser.parse_args(argv)
    args.log = setup_logging()

    # Get the daterange over which reduced files will be collected
    files = get_files(args)
    datestr = '%sv%03d' % (args.date, int(args.observation))
    filepath = '/work/00115/gebhardt/maverick/detect/%s/dithall.use' % datestr
    T = Table.read(filepath, format='ascii')

    # Creates a new file if the "--append" option is not set or the file
    # does not already exist.
    does_exist = False
    if op.exists(args.outfilename) and args.append:
        fileh = tb.open_file(args.outfilename, 'a')
        does_exist = True
    else:
        fileh = tb.open_file(args.outfilename, 'w')
        group = fileh.create_group(fileh.root, 'Info',
                                   'VIRUS Fiber Data and Metadata')
        table1 = fileh.create_table(group, 'Fibers', VIRUSFiber, 'Fiber Info')
        table2 = fileh.create_table(group, 'Shot', VIRUSShot, 'Shot Info')
        table3 = fileh.create_table(group, 'Images', VIRUSImage, 'Image Info')

    # Grab the fiber table and amplifier table for writing
    fibtable = fileh.root.Info.Fibers
    shottable = fileh.root.Info.Shot
    imagetable = fileh.root.Info.Images

    if does_exist:
        cnt = shottable[-1]['obsind']
    else:
        cnt = 1

    shot = shottable.row
    success = append_shot_to_table(shot, files[0], cnt)
    if success:
        shottable.flush()
    for fn in files:
        args.log.info('Working on %s' % fn)
        fib = fibtable.row
        im = imagetable.row
        success = append_fibers_to_table(fib, im, fn, cnt, T)
        if success:
            fibtable.flush()
            imagetable.flush()

    fileh.close()

if __name__ == '__main__':
    main()

max_stars_repo_name: grzeimann/Panacea

max_issues_repo_name: grzeimann/Panacea

max_forks_repo_name: grzeimann/Panacea

reason: import numpy,from astropy


text:
# small demo for sinogram TOF OS-MLEM

import os
import matplotlib.pyplot as plt
import pyparallelproj as ppp
from pyparallelproj.phantoms import ellipse2d_phantom, brain2d_phantom
from pyparallelproj.models import pet_fwd_model, pet_back_model

from scipy.ndimage import gaussian_filter
import numpy as np
import argparse

plt.ion()
plt.rcParams['image.cmap'] = 'Greys'

def osem(em_sino, attn_sino, sens_sino, contam_sino, proj, niter,
         fwhm = 0, verbose = False, xstart = None,
         callback = None, subset_callback = None,
         callback_kwargs = None, subset_callback_kwargs = None, vmax = None,
         fig = None, ax = None, figdir = None):

  img_shape  = tuple(proj.img_dim)

  # calculate the sensitivity images for each subset
  sens_img  = np.zeros((proj.nsubsets,) + img_shape, dtype = np.float32)

  for i in range(proj.nsubsets):
    # get the slice for the current subset
    ss        = proj.subset_slices[i]
    # generate a subset sinogram full of ones
    ones_sino = np.ones(proj.subset_sino_shapes[i] , dtype = np.float32)
    sens_img[i,...] = pet_back_model(ones_sino, proj, attn_sino[ss], sens_sino[ss], i, fwhm = fwhm)

  # initialize recon
  if xstart is None:
    recon = np.full(img_shape, (em_sino.sum()/sens_img.sum()), dtype = np.float32)
  else:
    recon = xstart.copy()

  if fig is not None:
    ax[0,3].imshow(em_sino.squeeze().T, aspect = 'auto', vmin = 0, vmax = em_sino.max())

  # run OSEM iterations
  for it in range(niter):
    for i in range(proj.nsubsets):
      if verbose: print(f'iteration {it+1} subset {i+1}')

      # get the slice for the current subset
      ss        = proj.subset_slices[i]

      exp_sino = pet_fwd_model(recon, proj, attn_sino[ss], sens_sino[ss], i,
                               fwhm = fwhm) + contam_sino[ss]
      ratio  = em_sino[ss] / exp_sino

      if fig is not None:
        if (it == 0) and (i == 0):
          p14 = ax[1,4].imshow(recon.squeeze(), vmin = 0, vmax = vmax)
        else:
          p14.set_data(recon.squeeze())

      back_ratio = pet_back_model(ratio, proj, attn_sino[ss], sens_sino[ss], i, fwhm = fwhm)

      recon *= (back_ratio / sens_img[i,...])

      if fig is not None:
        if (it == 0) and (i == 0):
          p13 = ax[1,3].imshow(exp_sino.squeeze().T, aspect = 'auto', vmin = 0, vmax = em_sino.max())
          p02 = ax[0,2].imshow(ratio.squeeze().T, aspect = 'auto', vmin = 0.8, vmax = 1.2)
          p01 = ax[0,1].imshow(back_ratio.squeeze(),
                               vmin = sens_img[i,...].min(),
                               vmax = sens_img[i,...].max())
          p11 = ax[1,1].imshow(sens_img[i,...].squeeze(),
                               vmin = sens_img[i,...].min(),
                               vmax = sens_img[i,...].max())
          p10 = ax[1,0].imshow(recon.squeeze(), vmin = 0, vmax = vmax)
          p00 = ax[0,0].imshow(back_ratio.squeeze()/sens_img[i,...].squeeze(), vmin = 0.5, vmax = 1.5)
          fig.tight_layout()
        else:
          p13.set_data(exp_sino.squeeze().T)
          p02.set_data(ratio.squeeze().T)
          p01.set_data(back_ratio.squeeze())
          p11.set_data(sens_img[i,...].squeeze())
          p00.set_data(back_ratio.squeeze()/sens_img[i,...].squeeze())
          p10.set_data(recon.squeeze())

        if (it < 10) or (it % 50) == 0:
          fig.savefig(os.path.join(figdir,f'iteration_{it+1}.png'))

      if subset_callback is not None:
        subset_callback(recon, iteration = (it+1), subset = (i+1), **subset_callback_kwargs)

    if callback is not None:
      callback(recon, iteration = (it+1), subset = (i+1), **callback_kwargs)

  return recon

#---------------------------------------------------------------------------------
# parse the command line

parser = argparse.ArgumentParser()
parser.add_argument('--ngpus',    help = 'number of GPUs to use', default = 0,   type = int)
parser.add_argument('--counts',   help = 'counts to simulate',    default = 1e6, type = float)
parser.add_argument('--niter',    help = 'number of iterations',  default = 4,   type = int)
parser.add_argument('--nsubsets',   help = 'number of subsets',     default = 28,  type = int)
parser.add_argument('--fwhm_mm',  help = 'psf modeling FWHM mm',  default = 4.5, type = float)
parser.add_argument('--fwhm_data_mm',  help = 'psf for data FWHM mm',  default = 4.5, type = float)
parser.add_argument('--phantom', help = 'phantom to use', default = 'brain2d')
parser.add_argument('--seed',    help = 'seed for random generator', default = 1, type = int)
args = parser.parse_args()

#---------------------------------------------------------------------------------

ngpus         = args.ngpus
counts        = args.counts
niter         = args.niter
nsubsets      = args.nsubsets
fwhm_mm       = args.fwhm_mm
fwhm_data_mm  = args.fwhm_data_mm
phantom       = args.phantom
seed          = args.seed

#---------------------------------------------------------------------------------

np.random.seed(seed)

# setup a scanner with one ring
scanner = ppp.RegularPolygonPETScanner(ncrystals_per_module = np.array([16,1]),
                                       nmodules             = np.array([28,1]))

# setup a test image
voxsize = np.array([2.,2.,2.])
n2      = max(1,int((scanner.xc2.max() - scanner.xc2.min()) / voxsize[2]))

# convert fwhm from mm to pixels
fwhm      = fwhm_mm / voxsize
fwhm_data = fwhm_data_mm / voxsize

# setup a test image
if phantom == 'ellipse2d':
  n   = 200
  img = np.zeros((n,n,n2), dtype = np.float32)
  tmp = ellipse_phantom(n = n, c = 3)
  for i2 in range(n2):
    img[:,:,i2] = tmp
elif phantom == 'brain2d':
  n   = 128
  img = np.zeros((n,n,n2), dtype = np.float32)
  tmp = brain2d_phantom(n = n)
  for i2 in range(n2):
    img[:,:,i2] = tmp

img_origin = (-(np.array(img.shape) / 2) +  0.5) * voxsize

# setup an attenuation image
att_img = (img > 0) * 0.01 * voxsize[0]

# generate nonTOF sinogram parameters and the nonTOF projector for attenuation projection
sino_params = ppp.PETSinogramParameters(scanner, rtrim = 146)
proj        = ppp.SinogramProjector(scanner, sino_params, img.shape, nsubsets = 1,
                                    voxsize = voxsize, img_origin = img_origin, ngpus = ngpus)

attn_sino = np.exp(-proj.fwd_project(att_img))

# generate the sensitivity sinogram
sens_sino = np.ones(sino_params.shape, dtype = np.float32)

# forward project the image
img_fwd= ppp.pet_fwd_model(img, proj, attn_sino, sens_sino, 0, fwhm = fwhm_data)

# scale sum of fwd image to counts
if counts > 0:
  scale_fac = (counts / img_fwd.sum())
  img_fwd  *= scale_fac
  img      *= scale_fac

  # contamination sinogram with scatter and randoms
  # useful to avoid division by 0 in the ratio of data and exprected data
  contam_sino = np.full(img_fwd.shape, 0.2*img_fwd.mean(), dtype = np.float32)

  em_sino = np.random.poisson(img_fwd + contam_sino)
else:
  scale_fac = 1.

  # contamination sinogram with sctter and randoms
  # useful to avoid division by 0 in the ratio of data and exprected data
  contam_sino = np.full(img_fwd.shape, 0.2*img_fwd.mean(), dtype = np.float32)

  em_sino = img_fwd + contam_sino

#-----------------------------------------------------------------------------------------------
# callback functions to calculate likelihood and show recon updates

def calc_cost(x):
  cost = 0

  for i in range(proj.nsubsets):
    # get the slice for the current subset
    ss = proj.subset_slices[i]
    exp = ppp.pet_fwd_model(x, proj, attn_sino[ss], sens_sino[ss], i, fwhm = fwhm) + contam_sino[ss]
    cost += (exp - em_sino[ss]*np.log(exp)).sum()

  return cost

def _cb(x, **kwargs):
  """ This function is called by the iterative recon algorithm after every iteration
      where x is the current reconstructed image
  """
  plt.pause(1e-6)
  it = kwargs.get('iteration',0)
  if 'cost' in kwargs:
    kwargs['cost'][it-1] = calc_cost(x)
  if 'intermed_recons' in kwargs:
    if it == 50:
      kwargs['intermed_recons'][0,...] = x
    if it == 100:
      kwargs['intermed_recons'][1,...] = x
    if it == 500:
      kwargs['intermed_recons'][2,...] = x
    if it == 1000:
      kwargs['intermed_recons'][3,...] = x

#-----------------------------------------------------------------------------------------------
# run the actual reconstruction using OSEM

# initialize the subsets for the projector
proj.init_subsets(nsubsets)

cost_osem = np.zeros(niter)

intermed_recons = np.zeros((4,) + img.shape)
cbk       = {'cost':cost_osem, 'intermed_recons':intermed_recons}

init_recon = np.full(img.shape, np.percentile(img,90), dtype = np.float32)

fig_osem, ax_osem = plt.subplots(2,5, figsize = (15,6))
for axx in ax_osem.ravel(): axx.set_axis_off()

figdir = f'mlem_iterations_{counts:.1e}'
if not os.path.exists(figdir): os.makedirs(figdir)

recon_osem = osem(em_sino, attn_sino, sens_sino, contam_sino, proj, niter,
                  fwhm = fwhm, verbose = True, xstart = init_recon, vmax = 1.2*img.max(),
                  callback = _cb, callback_kwargs = cbk)
                  #fig = fig_osem, ax = ax_osem, figdir = figdir)

#-----------------------------------------------------------------------------------------------
# plot the cost function

fig2, ax2 = plt.subplots(1,1, figsize = (5,3))
ax2.loglog(np.arange(1,niter+1), -cost_osem, '.-')
ax2.set_ylim(-cost_osem[3],-cost_osem[-1] + 0.1*(-cost_osem[-1]+cost_osem[2]))
ax2.set_ylabel('Poisson log-likelihood')
ax2.set_xlabel('iteration')
ax2.grid(ls = ':')
fig2.tight_layout()
fig2.savefig(os.path.join(figdir,'logL.png'))
fig2.show()

fig3, ax3 = plt.subplots(2,4, figsize = (11,6))
iters = [50,100,500,1000]
ps_fwhm_mm = 4.5
for i in range(4):
  ax3[0,i].imshow(intermed_recons[i,:,20:-20,0], vmin = 0, vmax = 1.2*img.max())
  ax3[1,i].imshow(gaussian_filter(intermed_recons[i,:,20:-20,0],ps_fwhm_mm/(2.35*voxsize[:2])),                  vmin = 0, vmax = 1.2*img.max())
  ax3[0,i].set_title(f'{iters[i]} updates', fontsize = 'small')
  ax3[1,i].set_title(f'{ps_fwhm_mm}mm smoothed', fontsize = 'small')
for axx in ax3.ravel():
  axx.set_axis_off()
fig3.tight_layout()
fig3.savefig(os.path.join(figdir,'conv.png'))
fig3.show()

max_stars_repo_name: gschramm/teaching

max_issues_repo_name: gschramm/teaching

max_forks_repo_name: gschramm/teaching

reason: import numpy,from scipy


text:
import csv
import numpy as np

from parameter import *

def get_similarity_table():
    with open(group_similarity_src, 'r', encoding='utf-8') as csvfile:
        spamreader = csv.reader(csvfile, delimiter=',')
        input_list = list(csv.reader(csvfile))

    index_array = input_list.pop(0)
    index_array.pop(0)   # pop the 0th empty field
    column_to_face_dict = dict()
    for key, face_id in enumerate(index_array):
        column_to_face_dict[key] = int(face_id)

    conversion_dict = dict()
    for conversion_array in input_list:
        this_face_id = conversion_array.pop(0)
        this_face_dict = dict()
        for key, rate in enumerate(conversion_array):
            this_face_dict[int(column_to_face_dict[key])] = float(rate)
        conversion_dict[int(this_face_id)] = this_face_dict

    return conversion_dict