<commit_before># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

from math import sin, cos, asin, pi, degrees, radians
import os

import numpy as np
import json

from .core import DiffractionPattern, DiffractionPatternCalculator, get_unique_families
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer

"""
This module implements a neutron diffraction (ND) pattern calculator.
"""

__author__ = "Yuta Suzuki"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "0.1"
__maintainer__ = "Yuta Suzuki"
__email__ = "resnant@outlook.jp"
__date__ = "4/19/18"

with open(os.path.join(os.path.dirname(__file__),
                       "neutron_scattering_length.json")) as f:
    # This table was cited from "Neutron Data Booklet" 2nd ed (Old City 2003).
    ATOMIC_SCATTERING_LEN = json.load(f)

class NMRPatternCalculator(DiffractionPatternCalculator):
    """
    Computes the powder neutron diffraction pattern of a crystal structure.
    This code is a slight modification of XRDCalculator in
    pymatgen.analysis.diffraction.xrd. See it for details of the algorithm.
    Main changes by using neutron instead of X-ray are as follows:
    1. Atomic scattering length is a constant.
    2. Polarization correction term of Lorentz factor is unnecessary.

    Reference:
    Marc De Graef and Michael E. McHenry, Structure of Materials 2nd ed,
    Chapter13, Cambridge University Press 2003.

    """

    # Tolerance in which to treat two peaks as having the same two theta.
    TWO_THETA_TOL = 1e-5

    # Tolerance in which to treat a peak as effectively 0 if the scaled
    # intensity is less than this number. Since the max intensity is 100,
    # this means the peak must be less than 1e-5 of the peak intensity to be
    # considered as zero. This deals with numerical issues where systematic
    # absences do not cancel exactly to zero.
    SCALED_INTENSITY_TOL = 1e-3

    def __init__(self, wavelength=1.54184, symprec=0,
                 debye_waller_factors=None):
        """
        Initializes the ND calculator with a given radiation.

        Args:
            wavelength (float): The wavelength of neutron in angstroms.
                Defaults to 1.54, corresponds to Cu K_alpha x-ray radiation.
            symprec (float): Symmetry precision for structure refinement. If
                set to 0, no refinement is done. Otherwise, refinement is
                performed using spglib with provided precision.
            debye_waller_factors ({element symbol: float}): Allows the
                specification of Debye-Waller factors. Note that these
                factors are temperature dependent.
        """
        self.wavelength = wavelength
        self.symprec = symprec
        self.debye_waller_factors = debye_waller_factors or {}

    def get_pattern(self, structure, scaled=True, two_theta_range=(0, 90)):
        """
        Calculates the powder neutron diffraction pattern for a structure.

        Args:
            structure (Structure): Input structure
            scaled (bool): Whether to return scaled intensities. The maximum
                peak is set to a value of 100. Defaults to True. Use False if
                you need the absolute values to combine ND plots.
            two_theta_range ([float of length 2]): Tuple for range of
                two_thetas to calculate in degrees. Defaults to (0, 90). Set to
                None if you want all diffracted beams within the limiting
                sphere of radius 2 / wavelength.

        Returns:
            (NDPattern)
        """
        if self.symprec:
            finder = SpacegroupAnalyzer(structure, symprec=self.symprec)
            structure = finder.get_refined_structure()

        wavelength = self.wavelength
        latt = structure.lattice
        is_hex = latt.is_hexagonal()

        # Obtained from Bragg condition. Note that reciprocal lattice
        # vector length is 1 / d_hkl.
        min_r, max_r = (0, 2 / wavelength) if two_theta_range is None else \
            [2 * sin(radians(t / 2)) / wavelength for t in two_theta_range]

        # Obtain crystallographic reciprocal lattice points within range
        recip_latt = latt.reciprocal_lattice_crystallographic
        recip_pts = recip_latt.get_points_in_sphere(
            [[0, 0, 0]], [0, 0, 0], max_r)
        if min_r:
            recip_pts = [pt for pt in recip_pts if pt[1] >= min_r]

        # Create a flattened array of coeffs, fcoords and occus. This is
        # used to perform vectorized computation of atomic scattering factors
        # later. Note that these are not necessarily the same size as the
        # structure as each partially occupied specie occupies its own
        # position in the flattened array.
        coeffs = []
        fcoords = []
        occus = []
        dwfactors = []

        for site in structure:
            for sp, occu in site.species_and_occu.items():
                try:
                    c = ATOMIC_SCATTERING_LEN[sp.symbol]
                except KeyError:
                    raise ValueError("Unable to calculate ND pattern as "
                                     "there is no scattering coefficients for"
                                     " %s." % sp.symbol)
                coeffs.append(c)
                dwfactors.append(self.debye_waller_factors.get(sp.symbol, 0))
                fcoords.append(site.frac_coords)
                occus.append(occu)

        coeffs = np.array(coeffs)
        fcoords = np.array(fcoords)
        occus = np.array(occus)
        dwfactors = np.array(dwfactors)
        peaks = {}
        two_thetas = []

        for hkl, g_hkl, ind in sorted(
                recip_pts, key=lambda i: (i[1], -i[0][0], -i[0][1], -i[0][2])):
            # Force miller indices to be integers.
            hkl = [int(round(i)) for i in hkl]
            if g_hkl != 0:

                d_hkl = 1 / g_hkl

                # Bragg condition
                theta = asin(wavelength * g_hkl / 2)

                # s = sin(theta) / wavelength = 1 / 2d = |ghkl| / 2 (d =
                # 1/|ghkl|)
                s = g_hkl / 2

                # Calculate Debye-Waller factor
                dw_correction = np.exp(-dwfactors * (s**2))

                # Vectorized computation of g.r for all fractional coords and
                # hkl.
                g_dot_r = np.dot(fcoords, np.transpose([hkl])).T[0]

                # Structure factor = sum of atomic scattering factors (with
                # position factor exp(2j * pi * g.r and occupancies).
                # Vectorized computation.
                f_hkl = np.sum(coeffs * occus * np.exp(2j * pi * g_dot_r)
                               * dw_correction)

                # Lorentz polarization correction for hkl
                lorentz_factor = 1 / (sin(theta) ** 2 * cos(theta))

                # Intensity for hkl is modulus square of structure factor.
                i_hkl = (f_hkl * f_hkl.conjugate()).real

                two_theta = degrees(2 * theta)

                if is_hex:
                    # Use Miller-Bravais indices for hexagonal lattices.
                    hkl = (hkl[0], hkl[1], - hkl[0] - hkl[1], hkl[2])
                # Deal with floating point precision issues.
                ind = np.where(np.abs(np.subtract(two_thetas, two_theta)) <
                               self.TWO_THETA_TOL)
                if len(ind[0]) > 0:
                    peaks[two_thetas[ind[0][0]]][0] += i_hkl * lorentz_factor
                    peaks[two_thetas[ind[0][0]]][1].append(tuple(hkl))
                else:
                    peaks[two_theta] = [i_hkl * lorentz_factor, [tuple(hkl)],
                                        d_hkl]
                    two_thetas.append(two_theta)

        # Scale intensities so that the max intensity is 100.
        max_intensity = max([v[0] for v in peaks.values()])
        x = []
        y = []
        hkls = []
        d_hkls = []
        for k in sorted(peaks.keys()):
            v = peaks[k]
            fam = get_unique_families(v[1])
            if v[0] / max_intensity * 100 > self.SCALED_INTENSITY_TOL:
                x.append(k)
                y.append(v[0])
                hkls.append(fam)
                d_hkls.append(v[2])
        nd = DiffractionPattern(x, y, hkls, d_hkls)
        if scaled:
            nd.normalize(mode="max", value=100)
        return nd

<commit_msg>More cleanup.<commit_after># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

from math import sin, cos, asin, pi, degrees, radians
import os

import numpy as np
import json

from .core import DiffractionPattern, DiffractionPatternCalculator, \
    get_unique_families
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer

"""
This module implements a neutron diffraction (ND) pattern calculator.
"""

__author__ = "Yuta Suzuki"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "0.1"
__maintainer__ = "Yuta Suzuki"
__email__ = "resnant@outlook.jp"
__date__ = "4/19/18"

with open(os.path.join(os.path.dirname(__file__),
                       "neutron_scattering_length.json")) as f:
    # This table was cited from "Neutron Data Booklet" 2nd ed (Old City 2003).
    ATOMIC_SCATTERING_LEN = json.load(f)

class NMRPatternCalculator(DiffractionPatternCalculator):
    """
    Computes the powder neutron diffraction pattern of a crystal structure.
    This code is a slight modification of XRDCalculator in
    pymatgen.analysis.diffraction.xrd. See it for details of the algorithm.
    Main changes by using neutron instead of X-ray are as follows:
    1. Atomic scattering length is a constant.
    2. Polarization correction term of Lorentz factor is unnecessary.

    Reference:
    Marc De Graef and Michael E. McHenry, Structure of Materials 2nd ed,
    Chapter13, Cambridge University Press 2003.

    """

    def __init__(self, wavelength=1.54184, symprec=0,
                 debye_waller_factors=None):
        """
        Initializes the ND calculator with a given radiation.

        Args:
            wavelength (float): The wavelength of neutron in angstroms.
                Defaults to 1.54, corresponds to Cu K_alpha x-ray radiation.
            symprec (float): Symmetry precision for structure refinement. If
                set to 0, no refinement is done. Otherwise, refinement is
                performed using spglib with provided precision.
            debye_waller_factors ({element symbol: float}): Allows the
                specification of Debye-Waller factors. Note that these
                factors are temperature dependent.
        """
        self.wavelength = wavelength
        self.symprec = symprec
        self.debye_waller_factors = debye_waller_factors or {}

    def get_pattern(self, structure, scaled=True, two_theta_range=(0, 90)):
        """
        Calculates the powder neutron diffraction pattern for a structure.

        Args:
            structure (Structure): Input structure
            scaled (bool): Whether to return scaled intensities. The maximum
                peak is set to a value of 100. Defaults to True. Use False if
                you need the absolute values to combine ND plots.
            two_theta_range ([float of length 2]): Tuple for range of
                two_thetas to calculate in degrees. Defaults to (0, 90). Set to
                None if you want all diffracted beams within the limiting
                sphere of radius 2 / wavelength.

        Returns:
            (NDPattern)
        """
        if self.symprec:
            finder = SpacegroupAnalyzer(structure, symprec=self.symprec)
            structure = finder.get_refined_structure()

        wavelength = self.wavelength
        latt = structure.lattice
        is_hex = latt.is_hexagonal()

        # Obtained from Bragg condition. Note that reciprocal lattice
        # vector length is 1 / d_hkl.
        min_r, max_r = (0, 2 / wavelength) if two_theta_range is None else \
            [2 * sin(radians(t / 2)) / wavelength for t in two_theta_range]

        # Obtain crystallographic reciprocal lattice points within range
        recip_latt = latt.reciprocal_lattice_crystallographic
        recip_pts = recip_latt.get_points_in_sphere(
            [[0, 0, 0]], [0, 0, 0], max_r)
        if min_r:
            recip_pts = [pt for pt in recip_pts if pt[1] >= min_r]

        # Create a flattened array of coeffs, fcoords and occus. This is
        # used to perform vectorized computation of atomic scattering factors
        # later. Note that these are not necessarily the same size as the
        # structure as each partially occupied specie occupies its own
        # position in the flattened array.
        coeffs = []
        fcoords = []
        occus = []
        dwfactors = []

        for site in structure:
            for sp, occu in site.species_and_occu.items():
                try:
                    c = ATOMIC_SCATTERING_LEN[sp.symbol]
                except KeyError:
                    raise ValueError("Unable to calculate ND pattern as "
                                     "there is no scattering coefficients for"
                                     " %s." % sp.symbol)
                coeffs.append(c)
                dwfactors.append(self.debye_waller_factors.get(sp.symbol, 0))
                fcoords.append(site.frac_coords)
                occus.append(occu)

        coeffs = np.array(coeffs)
        fcoords = np.array(fcoords)
        occus = np.array(occus)
        dwfactors = np.array(dwfactors)
        peaks = {}
        two_thetas = []

        for hkl, g_hkl, ind in sorted(
                recip_pts, key=lambda i: (i[1], -i[0][0], -i[0][1], -i[0][2])):
            # Force miller indices to be integers.
            hkl = [int(round(i)) for i in hkl]
            if g_hkl != 0:

                d_hkl = 1 / g_hkl

                # Bragg condition
                theta = asin(wavelength * g_hkl / 2)

                # s = sin(theta) / wavelength = 1 / 2d = |ghkl| / 2 (d =
                # 1/|ghkl|)
                s = g_hkl / 2

                # Calculate Debye-Waller factor
                dw_correction = np.exp(-dwfactors * (s**2))

                # Vectorized computation of g.r for all fractional coords and
                # hkl.
                g_dot_r = np.dot(fcoords, np.transpose([hkl])).T[0]

                # Structure factor = sum of atomic scattering factors (with
                # position factor exp(2j * pi * g.r and occupancies).
                # Vectorized computation.
                f_hkl = np.sum(coeffs * occus * np.exp(2j * pi * g_dot_r)
                               * dw_correction)

                # Lorentz polarization correction for hkl
                lorentz_factor = 1 / (sin(theta) ** 2 * cos(theta))

                # Intensity for hkl is modulus square of structure factor.
                i_hkl = (f_hkl * f_hkl.conjugate()).real

                two_theta = degrees(2 * theta)

                if is_hex:
                    # Use Miller-Bravais indices for hexagonal lattices.
                    hkl = (hkl[0], hkl[1], - hkl[0] - hkl[1], hkl[2])
                # Deal with floating point precision issues.
                ind = np.where(np.abs(np.subtract(two_thetas, two_theta)) <
                               self.TWO_THETA_TOL)
                if len(ind[0]) > 0:
                    peaks[two_thetas[ind[0][0]]][0] += i_hkl * lorentz_factor
                    peaks[two_thetas[ind[0][0]]][1].append(tuple(hkl))
                else:
                    peaks[two_theta] = [i_hkl * lorentz_factor, [tuple(hkl)],
                                        d_hkl]
                    two_thetas.append(two_theta)

        # Scale intensities so that the max intensity is 100.
        max_intensity = max([v[0] for v in peaks.values()])
        x = []
        y = []
        hkls = []
        d_hkls = []
        for k in sorted(peaks.keys()):
            v = peaks[k]
            fam = get_unique_families(v[1])
            if v[0] / max_intensity * 100 > self.SCALED_INTENSITY_TOL:
                x.append(k)
                y.append(v[0])
                hkls.append(fam)
                d_hkls.append(v[2])
        nd = DiffractionPattern(x, y, hkls, d_hkls)
        if scaled:
            nd.normalize(mode="max", value=100)
        return nd

<|endoftext|>


<commit_before>
import timetracker

def main(x):

    print('Starting campaign 002-iw')
    tt = timetracker.TimeTracker("002-iw")
    print(tt.get_campaign_runtime()[0])

    #iw_value_list = [4, 10, 20]
    #print('Running iw-value tests for iw-values {}'.format(str(iw_value_list)))

    x.exec('001-alive alpha')
    x.exec('001-alive beta')
    x.exec('001-alive koppa')

    # restore default sysctls and save it when initial
    x.exec('002-save-sysctls alpha')
    x.exec('002-save-sysctls beta')
    x.exec('002-save-sysctls koppa')
    x.exec('003-restore-sysctls alpha')
    x.exec('003-restore-sysctls beta')
    x.exec('003-restore-sysctls koppa')

    x.exec('105-offloading alpha offloading:off')
    x.exec('105-offloading beta offloading:off')

    x.exec('102-tcp-route-metrics alpha route-metrics-save:disabled')
    x.exec('102-tcp-route-metrics beta route-metrics-save:disabled')
    x.exec('102-tcp-route-metrics koppa route-metrics-save:disabled')

    print('##############################################')
    print('###  Starting iw[1-20] test without delay  ###')
    print('##############################################')
    # test iw[1-20] without network delay
    for current_value in range(1, 21):
        iw_value = str(current_value)
        print('Current iw-value tested: {}'.format(iw_value))

        # stop running interfering programs and delete possible qdiscs
        x.exec('007-kill alpha')
        x.exec('007-kill beta')
        x.exec('007-kill koppa')

        x.exec('009-network setup:dumbbell alpha beta')

        #x.exec('104-netem-cmd koppa control:partial change:add to-network:red command:"rate 10kbit"')
        #x.exec('104-netem-cmd koppa control:partial change:add to-network:blue command:"rate 10kbit"')

        x.exec('101-tcp-iw alpha initcwnd:{} initrwnd:20'.format(iw_value))
        x.exec('101-tcp-iw beta initcwnd:{} initrwnd:20'.format(iw_value))

        x.exec('006-tcpdump beta id:20 mode:start filter:"tcp and port 30000"')

        # test transfer
        x.exec('005-netperf alpha sink:beta id:21 source_port:20000 sink_port:30000 flow_length:-30000 flow_offset:1 netserver:29999')

        x.exec('008-wait-for-id-completion interval_time:2 alpha:21')

        x.exec('006-tcpdump beta id:20 mode:stop local-file-name:"./campaigns/002-iw/dumps/iw_{}.pcap"'.format(iw_value))

        print(tt.get_elapsed_runtime()[0])
        print(tt.get_remaining_runtime()[0])

    print('################################################')
    print('###  Starting iw[1-20] test with 20ms delay  ###')
    print('################################################')
    # test iw[1-20] with 20ms network delay
    for current_value in range(1, 21):
        iw_value = str(current_value)
        #for iw_value_list_position in range(0, len(iw_value_list)):
        #iw_value = str(iw_value_list[iw_value_list_position])
        print('Current iw-value tested: {}'.format(iw_value))

        # stop running interfering programs and delete possible qdiscs
        x.exec('007-kill alpha')
        x.exec('007-kill beta')
        x.exec('007-kill koppa')

        x.exec('009-network setup:dumbbell alpha beta')

        x.exec('104-netem-cmd koppa control:partial change:add to_network:red command:"delay 10ms"')
        x.exec('104-netem-cmd koppa control:partial change:add to_network:blue command:"delay 10ms"')

        #x.exec('104-netem-cmd koppa control:partial change:add to-network:red command:"rate 10kbit"')
        #x.exec('104-netem-cmd koppa control:partial change:add to-network:blue command:"rate 10kbit"')

        x.exec('101-tcp-iw alpha initcwnd:{} initrwnd:20'.format(iw_value))
        x.exec('101-tcp-iw beta initcwnd:{} initrwnd:20'.format(iw_value))

        x.exec('006-tcpdump beta id:20 mode:start filter:"tcp and port 30000"')

        # test transfer
        x.exec('005-netperf alpha sink:beta id:21 source_port:20000 sink_port:30000 flow_length:-30000 flow_offset:1 netserver:29999')

        x.exec('008-wait-for-id-completion interval_time:2 alpha:21')

        x.exec('006-tcpdump beta id:20 mode:stop local-file-name:"./campaigns/002-iw/dumps/iw_{}_delay_20ms.pcap"'.format(iw_value))

        print(tt.get_elapsed_runtime()[0])
        print(tt.get_remaining_runtime()[0])

    # reset qdiscs
    x.exec('104-netem-cmd koppa control:partial change:del to_network:red')
    x.exec('104-netem-cmd koppa control:partial change:del to_network:blue')

    # reset offloading
    x.exec('105-offloading alpha offloading:on')
    x.exec('105-offloading beta offloading:on')

    print(tt.update_campaign_runtime()[0])
<commit_msg>correcting campaign<commit_after>
import timetracker

def main(x):

    print('Starting campaign 002-iw')
    tt = timetracker.TimeTracker("002-iw")
    print(tt.get_campaign_runtime()[0])

    x.exec('001-alive alpha')
    x.exec('001-alive beta')
    x.exec('001-alive koppa')

    # restore default sysctls and save it when initial
    x.exec('002-save-sysctls alpha')
    x.exec('002-save-sysctls beta')
    x.exec('002-save-sysctls koppa')
    x.exec('003-restore-sysctls alpha')
    x.exec('003-restore-sysctls beta')
    x.exec('003-restore-sysctls koppa')

    x.exec('105-offloading alpha offloading:off')
    x.exec('105-offloading beta offloading:off')

    x.exec('102-tcp-route-metrics alpha route-metrics-save:disabled')
    x.exec('102-tcp-route-metrics beta route-metrics-save:disabled')
    x.exec('102-tcp-route-metrics koppa route-metrics-save:disabled')

    print('##############################################')
    print('###  Starting iw[1-20] test without delay  ###')
    print('##############################################')

    # test iw[1-20] without network delay
    for current_value in range(1, 21):
        iw_value = str(current_value)
        print('Current iw-value tested: {}'.format(iw_value))

        # stop running interfering programs and delete possible qdiscs
        x.exec('007-kill alpha')
        x.exec('007-kill beta')
        x.exec('007-kill koppa')

        x.exec('009-network setup:dumbbell alpha beta')

        #x.exec('104-netem-cmd koppa control:partial change:add to-network:red command:"rate 10kbit"')
        #x.exec('104-netem-cmd koppa control:partial change:add to-network:blue command:"rate 10kbit"')

        x.exec('101-tcp-iw alpha initcwnd:{} initrwnd:20'.format(iw_value))
        x.exec('101-tcp-iw beta initcwnd:{} initrwnd:20'.format(iw_value))

        x.exec('006-tcpdump alpha id:20 mode:start filter:"tcp and port 30000"')

        # test transfer
        x.exec('005-netperf alpha sink:beta id:21 source_port:20000 sink_port:30000 flow_length:-30000 flow_offset:1 netserver:29999')

        x.exec('008-wait-for-id-completion interval_time:2 alpha:21')

        x.exec('006-tcpdump alpha id:20 mode:stop local-file-name:"./campaigns/002-iw/dumps/iw_{}.pcap"'.format(iw_value))

        print(tt.get_elapsed_runtime()[0])
        print(tt.get_remaining_runtime()[0])

    print(tt.add_poi("iw_test_without_delay_done"))

    print('################################################')
    print('###  Starting iw[1-20] test with 20ms delay  ###')
    print('################################################')

    # test iw[1-20] with 20ms network delay
    for current_value in range(1, 21):
        iw_value = str(current_value)
        print('Current iw-value tested: {}'.format(iw_value))

        # stop running interfering programs and delete possible qdiscs
        x.exec('007-kill alpha')
        x.exec('007-kill beta')
        x.exec('007-kill koppa')

        x.exec('009-network setup:dumbbell alpha beta')

        x.exec('104-netem-cmd koppa control:partial change:add to_network:red command:"delay 10ms"')
        x.exec('104-netem-cmd koppa control:partial change:add to_network:blue command:"delay 10ms"')

        #x.exec('104-netem-cmd koppa control:partial change:add to-network:red command:"rate 10kbit"')
        #x.exec('104-netem-cmd koppa control:partial change:add to-network:blue command:"rate 10kbit"')

        x.exec('101-tcp-iw alpha initcwnd:{} initrwnd:20'.format(iw_value))
        x.exec('101-tcp-iw beta initcwnd:{} initrwnd:20'.format(iw_value))

        x.exec('006-tcpdump alpha id:20 mode:start filter:"tcp and port 30000"')

        # test transfer
        x.exec('005-netperf alpha sink:beta id:21 source_port:20000 sink_port:30000 flow_length:-30000 flow_offset:1 netserver:29999')

        x.exec('008-wait-for-id-completion interval_time:2 alpha:21')

        x.exec('006-tcpdump alpha id:20 mode:stop local-file-name:"./campaigns/002-iw/dumps/iw_{}_delay_20ms.pcap"'.format(iw_value))

        print(tt.get_elapsed_runtime()[0])
        print(tt.get_remaining_runtime()[0])

    # reset qdiscs
    x.exec('104-netem-cmd koppa control:partial change:del to_network:red')
    x.exec('104-netem-cmd koppa control:partial change:del to_network:blue')

    # reset offloading
    x.exec('105-offloading alpha offloading:on')
    x.exec('105-offloading beta offloading:on')

    print(tt.update_campaign_runtime()[0])
<|endoftext|>


<commit_before>import logging
import ibmsecurity.utilities.tools
from ibmsecurity.appliance.ibmappliance import IBMError

logger = logging.getLogger(__name__)

def get_all(isamAppliance, check_mode=False, force=False):
    """
    Retrieving a list of all CI server connections
    """
    return isamAppliance.invoke_get("Retrieving a list of all CI server connections",
                                    "/mga/server_connections/ci/v1")

def get(isamAppliance, name=None, check_mode=False, force=False):
    """
    Retrieving a CI server connection
    """
    ret_obj = _get_id(isamAppliance, name=name)
    id = ret_obj['data']

    if id == {}:
        return isamAppliance.create_return_object(1)
    else:
        return isamAppliance.invoke_get("Retrieving a CI server connection",
                                        "/mga/server_connections/ci/{0}/v1".format(id))

def get(isamAppliance, id=None, check_mode=False, force=False):
    """
    Retrieving a CI server connection
    """
    try:
        response = isamAppliance.invoke_get("Retrieving a CI server connection",
            "/mga/server_connections/ci/{0}/v1".format(id));
    except IBMError as e:
        if "404" in e[0]:
            response = isamAppliance.create_return_object(rc=404, data=e[1]);
        else:
            raise;

    return response;

def set(isamAppliance, name, connection, description='', locked=False, check_mode=False, force=False):
    """
    Creating or Modifying a CI server connection
    """
    if _check_exists(isamAppliance, name=name) is False:
        # Force the add - we already know connection does not exist
        return add(isamAppliance, name, connection, description, locked, check_mode, True)
    else:
        # Update request
        return update(isamAppliance, connection, description, locked, name, None,
                      check_mode, force)

def add(isamAppliance, name, connection, description='', locked=False, check_mode=False, force=False):
    """
    Creating a CI server connection
    """

    response = isamAppliance.create_return_object();

    if force is True or _check_exists(isamAppliance, name=name) is False:
        if check_mode is True:
            return isamAppliance.create_return_object(changed=True)
        else:
            try:
                response = isamAppliance.invoke_post(
                    "Creating a CI server connection",
                    "/mga/server_connections/ci/v1",
                    _create_json(name=name, description=description, locked=locked, connection=connection))

            except IBMError as e:
                if "400" in e[0]:
                    response = isamAppliance.create_return_object(rc=400, data=e[1]);
                elif "409" in e[0]:
                    response = isamAppliance.create_return_object(rc=409, data=e[1]);
                else:
                    raise;

    return response;

def deleteByName(isamAppliance, name=None, check_mode=False, force=False):
    """
    Deleting a CI server connection
    """
    if force is True or _check_exists(isamAppliance, name=name) is True:
        if check_mode is True:
            return isamAppliance.create_return_object(changed=True)
        else:
            ret_obj = _get_id(isamAppliance, name=name)
            id = ret_obj['data']
            return isamAppliance.invoke_delete(
                "Deleting a CI server connection",
                "/mga/server_connections/ci/{0}/v1".format(id))

    return isamAppliance.create_return_object()

def delete(isamAppliance, id, check_mode=False, force=False):
    """
    Deleting a CI server connection
    """
    if force is True or _check_exists(isamAppliance, id=id) is True:
        if check_mode is True:
            return isamAppliance.create_return_object(changed=True)
        else:
            try:
                response = isamAppliance.invoke_delete(
                    "Deleting a CI server connection",
                    "/mga/server_connections/ci/{0}/v1".format(id))
            except IBMError as e:
                if "404" in e[0]:
                    response = isamAppliance.create_return_object(rc=404, data=e[1]);
                else:
                    raise;

    return response;

def update(isamAppliance, id, connection, description='', locked=False, name=None,
           new_name=None, check_mode=False, force=False):
    """
    Modifying a CI server connection

    Use new_name to rename the connection, cannot compare password so update will take place everytime
    """
    if check_mode is True:
        return isamAppliance.create_return_object(changed=True)
    else:
        json_data = _create_json(name=name, description=description, locked=locked, connection=connection)
        if new_name is not None:  # Rename condition
            json_data['name'] = new_name
        try:
            response = isamAppliance.invoke_put(
                "Modifying a CI server connection",
                "/mga/server_connections/ci/{0}/v1".format(id), json_data)
        except IBMError as e:
            if "400" in e[0]:
                response = isamAppliance.create_return_object(rc=400, data=e[1]);
            else:
                raise;

        return response;

def _create_json(name, description, locked, connection):
    """
    Create a JSON to be used for the REST API call
    """
    json = {
        "connection": connection,
        "type": "ci",
        "name": name,
        "description": description,
        "locked": locked
    }

    return json

def _get_id(isamAppliance, name):
    """
    Retrieve UUID for named CI connection
    """
    ret_obj = get_all(isamAppliance)

    ret_obj_new = isamAppliance.create_return_object()

    for obj in ret_obj['data']:
        if obj['name'] == name:
            ret_obj_new['data'] = obj['uuid']

    return ret_obj_new

def _check_exists(isamAppliance, name=None, id=None):
    """
    Check if CI Connection already exists
    """
    ret_obj = get_all(isamAppliance)

    for obj in ret_obj['data']:
        if (name is not None and obj['name'] == name) or (id is not None and obj['uuid'] == id):
            return True

    return False

def compare(isamAppliance1, isamAppliance2):
    """
    Compare CI Connections between two appliances
    """
    ret_obj1 = get_all(isamAppliance1)
    ret_obj2 = get_all(isamAppliance2)

    for obj in ret_obj1['data']:
        del obj['uuid']
    for obj in ret_obj2['data']:
        del obj['uuid']

    return ibmsecurity.utilities.tools.json_compare(ret_obj1, ret_obj2, deleted_keys=['uuid'])
<commit_msg>Make CI Server Connections API consistent with other Server Connections APIs<commit_after>import logging
import ibmsecurity.utilities.tools
from ibmsecurity.appliance.ibmappliance import IBMError

logger = logging.getLogger(__name__)

def get_all(isamAppliance, check_mode=False, force=False):
    """
    Retrieving a list of all CI server connections
    """
    return isamAppliance.invoke_get("Retrieving a list of all CI server connections",
                                    "/mga/server_connections/ci/v1")

def get(isamAppliance, name, check_mode=False, force=False):
    """
    Retrieving a CI server connection
    """
    ret_obj = _get_id(isamAppliance, name=name)
    id = ret_obj['data']

    if id == {}:
        return isamAppliance.create_return_object(1)
    else:
        return isamAppliance.invoke_get("Retrieving a CI server connection",
                                        "/mga/server_connections/ci/{0}/v1".format(id))

def set(isamAppliance, name, connection, description='', locked=False, check_mode=False, force=False):
    """
    Creating or Modifying a CI server connection
    """
    if _check_exists(isamAppliance, name=name) is False:
        # Force the add - we already know connection does not exist
        return add(isamAppliance, name, connection, description, locked, check_mode, True)
    else:
        # Update request
        return update(isamAppliance, connection, description, locked, name, None,
                      check_mode, force)

def add(isamAppliance, name, connection, description='', locked=False, check_mode=False, force=False):
    """
    Creating a CI server connection
    """

    response = isamAppliance.create_return_object();

    if force is True or _check_exists(isamAppliance, name=name) is False:
        if check_mode is True:
            return isamAppliance.create_return_object(changed=True)
        else:
            try:
                response = isamAppliance.invoke_post(
                    "Creating a CI server connection",
                    "/mga/server_connections/ci/v1",
                    _create_json(name=name, description=description, locked=locked, connection=connection))

            except IBMError as e:
                if "400" in e[0]:
                    response = isamAppliance.create_return_object(rc=400, data=e[1]);
                elif "409" in e[0]:
                    response = isamAppliance.create_return_object(rc=409, data=e[1]);
                else:
                    raise;

    return response;

def deleteByName(isamAppliance, name=None, check_mode=False, force=False):
    """
    Deleting a CI server connection
    """
    if force is True or _check_exists(isamAppliance, name=name) is True:
        if check_mode is True:
            return isamAppliance.create_return_object(changed=True)
        else:
            ret_obj = _get_id(isamAppliance, name=name)
            id = ret_obj['data']
            return isamAppliance.invoke_delete(
                "Deleting a CI server connection",
                "/mga/server_connections/ci/{0}/v1".format(id))

    return isamAppliance.create_return_object()

def delete(isamAppliance, id, check_mode=False, force=False):
    """
    Deleting a CI server connection
    """
    if force is True or _check_exists(isamAppliance, id=id) is True:
        if check_mode is True:
            return isamAppliance.create_return_object(changed=True)
        else:
            try:
                response = isamAppliance.invoke_delete(
                    "Deleting a CI server connection",
                    "/mga/server_connections/ci/{0}/v1".format(id))
            except IBMError as e:
                if "404" in e[0]:
                    response = isamAppliance.create_return_object(rc=404, data=e[1]);
                else:
                    raise;

    return response;

def update(isamAppliance, id, connection, description='', locked=False, name=None,
           new_name=None, check_mode=False, force=False):
    """
    Modifying a CI server connection

    Use new_name to rename the connection, cannot compare password so update will take place everytime
    """
    if check_mode is True:
        return isamAppliance.create_return_object(changed=True)
    else:
        json_data = _create_json(name=name, description=description, locked=locked, connection=connection)
        if new_name is not None:  # Rename condition
            json_data['name'] = new_name
        try:
            response = isamAppliance.invoke_put(
                "Modifying a CI server connection",
                "/mga/server_connections/ci/{0}/v1".format(id), json_data)
        except IBMError as e:
            if "400" in e[0]:
                response = isamAppliance.create_return_object(rc=400, data=e[1]);
            else:
                raise;

        return response;

def _create_json(name, description, locked, connection):
    """
    Create a JSON to be used for the REST API call
    """
    json = {
        "connection": connection,
        "type": "ci",
        "name": name,
        "description": description,
        "locked": locked
    }

    return json

def _get_id(isamAppliance, name):
    """
    Retrieve UUID for named CI connection
    """
    ret_obj = get_all(isamAppliance)

    ret_obj_new = isamAppliance.create_return_object()

    for obj in ret_obj['data']:
        if obj['name'] == name:
            ret_obj_new['data'] = obj['uuid']

    return ret_obj_new

def _check_exists(isamAppliance, name=None, id=None):
    """
    Check if CI Connection already exists
    """
    ret_obj = get_all(isamAppliance)

    for obj in ret_obj['data']:
        if (name is not None and obj['name'] == name) or (id is not None and obj['uuid'] == id):
            return True

    return False

def compare(isamAppliance1, isamAppliance2):
    """
    Compare CI Connections between two appliances
    """
    ret_obj1 = get_all(isamAppliance1)
    ret_obj2 = get_all(isamAppliance2)

    for obj in ret_obj1['data']:
        del obj['uuid']
    for obj in ret_obj2['data']:
        del obj['uuid']

    return ibmsecurity.utilities.tools.json_compare(ret_obj1, ret_obj2, deleted_keys=['uuid'])
<|endoftext|>


<commit_before>from django import forms
from django.utils.safestring import mark_safe
from django.forms.util import flatatt
from django.conf import settings

class SliderWidget(forms.TextInput):
    """
    http://pastebin.com/f34f0c71d
    """

    def __init__(self, min=None, max=None, step=None, image=None, attrs=None):
        super(SliderWidget, self).__init__(attrs)
        self.max = 100
        self.min = 0
        self.step = None
        self.image = None

        if max:
            self.max = max
        if min:
            self.min = min
        if step:
            self.step = step
        if image:
            self.image = image

    def get_step(self):
        if self.step:
            return "step : %s," % self.step
        else:
            return ''

    def render(self, name, value, attrs=None):
        attrs['class'] = 'slidervalue'
        final_attrs = self.build_attrs(attrs, name=name)
        slider_id = 'slider-'+name

        field = super(SliderWidget, self).render(name, value, attrs)
        image_html = ""
        if self.image:
            url = self.image
            if not url.startswith("/") and not self.image.startswith("http://"):
                url = settings.MEDIA_URL + url
            image_html = """""" % url
        slider = """

        """ % { 'slider_id' : slider_id,
                'field_id' : "id_%s" % name,
                'min' : self.min,
                'max' : self.max,
                'step' : self.get_step()}

        return mark_safe(image_html+field+slider)

<commit_msg>Slider widget to update values as it's sliding<commit_after>from django import forms
from django.utils.safestring import mark_safe
from django.forms.util import flatatt
from django.conf import settings

class SliderWidget(forms.TextInput):
    """
    http://pastebin.com/f34f0c71d
    """

    def __init__(self, min=None, max=None, step=None, image=None, attrs=None):
        super(SliderWidget, self).__init__(attrs)
        self.max = 100
        self.min = 0
        self.step = None
        self.image = None

        if max:
            self.max = max
        if min:
            self.min = min
        if step:
            self.step = step
        if image:
            self.image = image

    def get_step(self):
        if self.step:
            return "step : %s," % self.step
        else:
            return ''

    def render(self, name, value, attrs=None):
        attrs['class'] = 'slidervalue'
        final_attrs = self.build_attrs(attrs, name=name)
        slider_id = 'slider-'+name

        field = super(SliderWidget, self).render(name, value, attrs)
        image_html = ""
        if self.image:
            url = self.image
            if not url.startswith("/") and not self.image.startswith("http://"):
                url = settings.MEDIA_URL + url
            image_html = """""" % url
        slider = """

        """ % { 'slider_id' : slider_id,
                'field_id' : "id_%s" % name,
                'min' : self.min,
                'max' : self.max,
                'step' : self.get_step()}

        return mark_safe(image_html+field+slider)

<|endoftext|>


<commit_before>"""
token authentication module
"""

from django.conf import settings
from django.contrib.auth.models import User
from tardis.tardis_portal.models import Token, ObjectACL, Experiment
from tardis.tardis_portal.auth.interfaces import GroupProvider

TOKEN_EXPERIMENT = '_token_experiment'

def _ensure_acl_exists(experiment_id):
    experiment = Experiment.objects.get(pk=experiment_id)

    ObjectACL.objects.get_or_create(
        pluginId=TokenGroupProvider.name,
        entityId=str(experiment.id), canRead=True,
        content_type=experiment.get_ct(),
        object_id=experiment.id,
        aclOwnershipType=ObjectACL.OWNER_OWNED)

# def authenticate(request, token_string):
#     try:
#         token = Token.objects.get(token=token_string)
#     except Token.DoesNotExist:
#         return None
#     else:
#         if token.is_expired():
#             return None

#     user = User.objects.get(username=settings.TOKEN_USERNAME)
#     user.backend = 'django.contrib.auth.backends.ModelBackend'

#     request.session[TOKEN_EXPERIMENT] = token.experiment.id
#     _ensure_acl_exists(token.experiment.id)
#     request.session.set_expiry(token.get_session_expiry())

#     return user

class TokenGroupProvider(GroupProvider):
    '''
    '''
    name = u'token_group'

    def getGroups(self, user):
        if hasattr(user, 'allowed_tokens'):
            tokens = Token.objects.filter(
                token__in=user.allowed_tokens)
            experiment_ids = []
            for token in tokens:
                if not token.is_expired():
                    _ensure_acl_exists(token.experiment.id)
                    experiment_ids.append(token.experiment.id)
            return experiment_ids
        else:
            return []

    def searchGroups(self, **kwargs):
        """
            return nothing because these are not groups in
            the standard sense
        """
        return []

class TokenAuthMiddleware(object):
    '''
    adds tokens to the user object and the session from a GET query
    '''

    def process_request(self, request):
        if 'token' in request.GET:
            all_tokens = set()
            all_tokens.add(request.GET['token'])
            all_tokens.update(getattr(request.user, 'allowed_tokens', []))
            all_tokens.update(request.session.get('allowed_tokens', []))
            all_tokens = list(all_tokens)
            request.user.allowed_tokens = all_tokens
            request.session['allowed_tokens'] = all_tokens
        return None
<commit_msg>fixed bug with token auth<commit_after>"""
token authentication module
"""

from django.conf import settings
from django.contrib.auth.models import User
from tardis.tardis_portal.models import Token, ObjectACL, Experiment
from tardis.tardis_portal.auth.interfaces import GroupProvider

TOKEN_EXPERIMENT = '_token_experiment'

def _ensure_acl_exists(experiment_id):
    experiment = Experiment.objects.get(pk=experiment_id)

    ObjectACL.objects.get_or_create(
        pluginId=TokenGroupProvider.name,
        entityId=str(experiment.id), canRead=True,
        content_type=experiment.get_ct(),
        object_id=experiment.id,
        aclOwnershipType=ObjectACL.OWNER_OWNED)

# def authenticate(request, token_string):
#     try:
#         token = Token.objects.get(token=token_string)
#     except Token.DoesNotExist:
#         return None
#     else:
#         if token.is_expired():
#             return None

#     user = User.objects.get(username=settings.TOKEN_USERNAME)
#     user.backend = 'django.contrib.auth.backends.ModelBackend'

#     request.session[TOKEN_EXPERIMENT] = token.experiment.id
#     _ensure_acl_exists(token.experiment.id)
#     request.session.set_expiry(token.get_session_expiry())

#     return user

class TokenGroupProvider(GroupProvider):
    '''
    '''
    name = u'token_group'

    def getGroups(self, user):
        if hasattr(user, 'allowed_tokens'):
            tokens = Token.objects.filter(
                token__in=user.allowed_tokens)
            experiment_ids = []
            for token in tokens:
                if not token.is_expired():
                    _ensure_acl_exists(token.experiment.id)
                    experiment_ids.append(token.experiment.id)
            return experiment_ids
        else:
            return []

    def searchGroups(self, **kwargs):
        """
            return nothing because these are not groups in
            the standard sense
        """
        return []

class TokenAuthMiddleware(object):
    '''
    adds tokens to the user object and the session from a GET query
    '''

    def process_request(self, request):
        all_tokens = set()
        all_tokens.add(request.GET.get('token', None))
        all_tokens.update(getattr(request.user, 'allowed_tokens', []))
        all_tokens.update(request.session.get('allowed_tokens', []))
        all_tokens = list(all_tokens)
        request.user.allowed_tokens = all_tokens
        request.session['allowed_tokens'] = all_tokens
        return None
<|endoftext|>


<commit_before>import numpy as np

#################################################################
# Implements multi-agent controllers
#################################################################

class RandomPolicy():

    # constructor
    def __init__(self, n_actions, rng = np.random.RandomState()):
        self.rng = rng
        self.n_actions = n_actions

    def action(self, state):
        return self.rng.randint(self.n_actions)

<commit_msg>renamed policy file<commit_after><|endoftext|>


<commit_before>import argparse
from sklearn import cross_validation

from similarity_utils import *
from readfacedatabases import *
import similarity

parser = argparse.ArgumentParser(description='digit recognition')
parser.add_argument('--relu', dest='relu',action='store_true', default=False,
                    help=("if true, trains the RBM or DBN with a rectified linear unit"))
parser.add_argument('--cv', dest='cv',action='store_true', default=False,
                    help=("if true, does cv"))
parser.add_argument('--testKanadeMain', dest='testKanadeMain',action='store_true', default=False,
                    help=("if true, tests the net with the Kanade databse"))
parser.add_argument('--diffsubjects', dest='diffsubjects',action='store_true', default=False,
                    help=("if true, trains a net with different test and train subjects"))

args = parser.parse_args()

def similarityMain():
  trainData1, trainData2, testData1, testData2, similaritiesTrain, similaritiesTest = splitData()

  print "training with dataset of size ", len(train)
  print len(trainData1)

  print "testing with dataset of size ", len(test)
  print len(testData1)

  simNet = similarity.SimilarityNet(learningRate=0.001,
                                    maxMomentum=0.95,
                                    binary=True,
                                    rbmNrVis=1200,
                                    rbmNrHid=1000,
                                    rbmLearningRate=0.005,
                                    rbmDropoutHid=1.0,
                                    rbmDropoutVis=1.0,
                                    trainingEpochsRBM=10)

  simNet.train(trainData1, trainData2, similaritiesTrain)

  res = simNet.test(testData1, testData2)

  predicted = res > 0.5

  correct = (similaritiesTest == predicted).sum() * 1.0 / len(res)

  print res

  print correct

def similarityMainTestKanade():
  subjectsToImgs = readCroppedYaleSubjects()
  trainData1, trainData2, trainSubjects1, trainSubjects2 =\
    splitDataAccordingToSubjects(subjectsToImgs, None, imgsPerSubject=None)

  testData1, testData2, similaritiesTest = splitSimilarityKanade()

  similaritiesTrain =  similarityDifferentSubjects(trainSubjects1, trainSubjects2)

  print "training with dataset of size ", len(train)
  print len(trainData1)

  print "testing with dataset of size ", len(test)
  print len(testData1)

  simNet = similarity.SimilarityNet(learningRate=0.001,
                                    maxMomentum=0.95,
                                    binary=True,
                                    rbmNrVis=1200,
                                    rbmNrHid=1000,
                                    rbmLearningRate=0.005,
                                    rbmDropoutHid=1.0,
                                    rbmDropoutVis=1.0,
                                    trainingEpochsRBM=10)

  simNet.train(trainData1, trainData2, similaritiesTrain)

  res = simNet.test(testData1, testData2)

  predicted = res > 0.5

  correct = (similaritiesTest == predicted).sum() * 1.0 / len(res)

  print res

  print correct

def similarityDifferentSubjectsMain():
  nrSubjects = 147
  subjects = np.array(range(nrSubjects))
  kf = cross_validation.KFold(n=len(subjects), k=5)

  for train, test in kf:
    break

  subjectsToImgs = readMultiPIESubjects()

  subjectTrain = subjects[train]
  subjectTest = subjects[test]

  print "len(subjectTrain)"
  print len(subjectTrain)
  print "len(subjectTest)"
  print len(subjectTest)

  trainData1, trainData2, trainSubjects1, trainSubjects2 =\
    splitDataAccordingToSubjects(subjectsToImgs, subjectTrain, imgsPerSubject=None)

  testData1, testData2, testSubjects1, testSubjects2 =\
    splitDataAccordingToSubjects(subjectsToImgs, subjectTest, imgsPerSubject=None)

  print "training with dataset of size ", len(trainData1)
  print "testing with dataset of size ", len(testData1)

  similaritiesTrain =  similarityDifferentSubjects(trainSubjects1, trainSubjects2)
  similaritiesTest =  similarityDifferentSubjects(testSubjects1, testSubjects2)

  simNet = similarity.SimilarityNet(learningRate=0.001,
                                    maxMomentum=0.95,
                                    binary=True,
                                    rbmNrVis=1200,
                                    rbmNrHid=1000,
                                    rbmLearningRate=0.005,
                                    rbmDropoutHid=1.0,
                                    rbmDropoutVis=1.0,
                                    trainingEpochsRBM=10)

  simNet.train(trainData1, trainData2, similaritiesTrain)

  res = simNet.test(testData1, testData2)

  predicted = res > 0.5

  correct = (similaritiesTest == predicted).sum() * 1.0 / len(res)

  print res

  print correct

def similarityCV():
  trainData1, trainData2, testData1, testData2, similaritiesTrain, similaritiesTest = splitData()

  params = [(0.0001, 0.01), (0.0001, 0.005), (0.001, 0.01), (0.001, 0.005)]
  kf = cross_validation.KFold(n=len(trainData1), k=len(params))

  fold = 0
  for train, test in kf:
    simNet = similarity.SimilarityNet(learningRate=params[fold][0],
                                    maxMomentum=0.95,
                                    binary=True,
                                    rbmNrVis=1200,
                                    rbmNrHid=1000,
                                    rbmLearningRate=params[fold][1],
                                    rbmDropoutHid=1.0,
                                    rbmDropoutVis=1.0)

    simNet.train(trainData1, trainData2, similaritiesTrain)

    res = simNet.test(testData1, testData2)

    # error = (similaritiesTest - res)  * 1.0 / len(res)

    predicted = res > 0.5

    print "predicted"
    print predicted

    correct = (similaritiesTest == predicted).sum() * 1.0 / len(res)

    print "params[fold]"
    print params[fold]
    print "res"
    print res
    print "correct"
    print correct

    fold += 1

def main():
  if args.cv:
    similarityCV()
  if args.diffsubjects:
    similarityDifferentSubjectsMain()
  if args.testKanadeMain:
    similarityMainTestKanade()
  else:
    similarityMain()

if __name__ == '__main__':
  main()
<commit_msg>fix typo<commit_after>import argparse
from sklearn import cross_validation

from similarity_utils import *
from readfacedatabases import *
import similarity

parser = argparse.ArgumentParser(description='digit recognition')
parser.add_argument('--relu', dest='relu',action='store_true', default=False,
                    help=("if true, trains the RBM or DBN with a rectified linear unit"))
parser.add_argument('--cv', dest='cv',action='store_true', default=False,
                    help=("if true, does cv"))
parser.add_argument('--testKanadeMain', dest='testKanadeMain',action='store_true', default=False,
                    help=("if true, tests the net with the Kanade databse"))
parser.add_argument('--diffsubjects', dest='diffsubjects',action='store_true', default=False,
                    help=("if true, trains a net with different test and train subjects"))

args = parser.parse_args()

def similarityMain():
  trainData1, trainData2, testData1, testData2, similaritiesTrain, similaritiesTest = splitData()

  print "training with dataset of size ", len(train)
  print len(trainData1)

  print "testing with dataset of size ", len(test)
  print len(testData1)

  simNet = similarity.SimilarityNet(learningRate=0.001,
                                    maxMomentum=0.95,
                                    binary=True,
                                    rbmNrVis=1200,
                                    rbmNrHid=1000,
                                    rbmLearningRate=0.005,
                                    rbmDropoutHid=1.0,
                                    rbmDropoutVis=1.0,
                                    trainingEpochsRBM=10)

  simNet.train(trainData1, trainData2, similaritiesTrain)

  res = simNet.test(testData1, testData2)

  predicted = res > 0.5

  correct = (similaritiesTest == predicted).sum() * 1.0 / len(res)

  print res

  print correct

def similarityMainTestKanade():
  subjectsToImgs = readCroppedYaleSubjects()
  trainData1, trainData2, trainSubjects1, trainSubjects2 =\
    splitDataAccordingToSubjects(subjectsToImgs, None, imgsPerSubject=None)

  testData1, testData2, similaritiesTest = splitSimilarityKanade()

  similaritiesTrain =  similarityDifferentSubjects(trainSubjects1, trainSubjects2)

  print "training with dataset of size ", len(trainData1)
  print len(trainData1)

  print "testing with dataset of size ", len(testData1)
  print len(testData1)

  simNet = similarity.SimilarityNet(learningRate=0.001,
                                    maxMomentum=0.95,
                                    binary=True,
                                    rbmNrVis=1200,
                                    rbmNrHid=1000,
                                    rbmLearningRate=0.005,
                                    rbmDropoutHid=1.0,
                                    rbmDropoutVis=1.0,
                                    trainingEpochsRBM=10)

  simNet.train(trainData1, trainData2, similaritiesTrain)

  res = simNet.test(testData1, testData2)

  predicted = res > 0.5

  correct = (similaritiesTest == predicted).sum() * 1.0 / len(res)

  print res

  print correct

def similarityDifferentSubjectsMain():
  nrSubjects = 147
  subjects = np.array(range(nrSubjects))
  kf = cross_validation.KFold(n=len(subjects), k=5)

  for train, test in kf:
    break

  subjectsToImgs = readMultiPIESubjects()

  subjectTrain = subjects[train]
  subjectTest = subjects[test]

  print "len(subjectTrain)"
  print len(subjectTrain)
  print "len(subjectTest)"
  print len(subjectTest)

  trainData1, trainData2, trainSubjects1, trainSubjects2 =\
    splitDataAccordingToSubjects(subjectsToImgs, subjectTrain, imgsPerSubject=None)

  testData1, testData2, testSubjects1, testSubjects2 =\
    splitDataAccordingToSubjects(subjectsToImgs, subjectTest, imgsPerSubject=None)

  print "training with dataset of size ", len(trainData1)
  print "testing with dataset of size ", len(testData1)

  similaritiesTrain =  similarityDifferentSubjects(trainSubjects1, trainSubjects2)
  similaritiesTest =  similarityDifferentSubjects(testSubjects1, testSubjects2)

  simNet = similarity.SimilarityNet(learningRate=0.001,
                                    maxMomentum=0.95,
                                    binary=True,
                                    rbmNrVis=1200,
                                    rbmNrHid=1000,
                                    rbmLearningRate=0.005,
                                    rbmDropoutHid=1.0,
                                    rbmDropoutVis=1.0,
                                    trainingEpochsRBM=10)

  simNet.train(trainData1, trainData2, similaritiesTrain)

  res = simNet.test(testData1, testData2)

  predicted = res > 0.5

  correct = (similaritiesTest == predicted).sum() * 1.0 / len(res)

  print res

  print correct

def similarityCV():
  trainData1, trainData2, testData1, testData2, similaritiesTrain, similaritiesTest = splitData()

  params = [(0.0001, 0.01), (0.0001, 0.005), (0.001, 0.01), (0.001, 0.005)]
  kf = cross_validation.KFold(n=len(trainData1), k=len(params))

  fold = 0
  for train, test in kf:
    simNet = similarity.SimilarityNet(learningRate=params[fold][0],
                                    maxMomentum=0.95,
                                    binary=True,
                                    rbmNrVis=1200,
                                    rbmNrHid=1000,
                                    rbmLearningRate=params[fold][1],
                                    rbmDropoutHid=1.0,
                                    rbmDropoutVis=1.0)

    simNet.train(trainData1, trainData2, similaritiesTrain)

    res = simNet.test(testData1, testData2)

    # error = (similaritiesTest - res)  * 1.0 / len(res)

    predicted = res > 0.5

    print "predicted"
    print predicted

    correct = (similaritiesTest == predicted).sum() * 1.0 / len(res)

    print "params[fold]"
    print params[fold]
    print "res"
    print res
    print "correct"
    print correct

    fold += 1

def main():
  if args.cv:
    similarityCV()
  if args.diffsubjects:
    similarityDifferentSubjectsMain()
  if args.testKanadeMain:
    similarityMainTestKanade()
  else:
    similarityMain()

if __name__ == '__main__':
  main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

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

import django_filters

from . import models as emg_models

WORD_MATCH_REGEX = r"{0}"

class PublicationFilter(django_filters.FilterSet):

    data_origination = django_filters.CharFilter(
        name='studies__data_origination',
        distinct=True)

    class Meta:
        model = emg_models.Publication
        fields = (
            'data_origination',
        )

class StudyFilter(django_filters.FilterSet):

    biome = django_filters.CharFilter(method='filter_biome', distinct=True)

    def filter_biome(self, qs, name, value):
        try:
            b = emg_models.Biome.objects.get(lineage=value)
            studies = emg_models.Sample.objects.values('study') \
                .available(self.request) \
                .filter(
                    biome__lft__gte=b.lft-1,
                    biome__rgt__lte=b.rgt+1)
            qs = qs.filter(pk__in=studies)
        except emg_models.Biome.DoesNotExist:
            pass
        return qs

    other_accession = django_filters.CharFilter(
        name='project_id',
        distinct=True)

    centre_name = django_filters.CharFilter(
        method='filter_centre_name', distinct=True)

    def filter_centre_name(self, qs, name, value):
        return qs.filter(
            centre_name__iregex=WORD_MATCH_REGEX.format(value))

    class Meta:
        model = emg_models.Study
        fields = (
            'biome',
            'centre_name',
            'other_accession',
        )

class SampleFilter(django_filters.FilterSet):

    experiment_type = django_filters.CharFilter(
        name='runs__experiment_type__experiment_type',
        distinct=True)

    # pipeline_version = django_filters.CharFilter(
    #     name='runs__pipeline__release_version',
    #     distinct=True)

    biome = django_filters.CharFilter(method='filter_biome', distinct=True)

    def filter_biome(self, qs, name, value):
        try:
            b = emg_models.Biome.objects.get(lineage=value)
            qs = qs.filter(biome__lft__gte=b.lft-1, biome__rgt__lte=b.rgt+1)
        except emg_models.Biome.DoesNotExist:
            pass
        return qs

    instrument_platform = django_filters.CharFilter(
        method='filter_instrument_platform', distinct=True)

    def filter_instrument_platform(self, qs, name, value):
        samples = emg_models.Run.objects.values('sample_id') \
            .available(self.request) \
            .filter(instrument_platform__iregex=WORD_MATCH_REGEX.format(value))
        return qs.filter(pk__in=samples).select_related('study', 'biome')

    instrument_model = django_filters.CharFilter(
        method='filter_instrument_model', distinct=True)

    def filter_instrument_model(self, qs, name, value):
        samples = emg_models.Run.objects.values('sample_id') \
            .available(self.request) \
            .filter(instrument_model__iregex=WORD_MATCH_REGEX.format(value))
        return qs.filter(pk__in=samples).select_related('study', 'biome')

    metadata_key = django_filters.CharFilter(
        method='filter_metadata_key', distinct=True)

    def filter_metadata_key(self, qs, name, value):
        return qs.filter(
            metadata__var__var_name__iregex=WORD_MATCH_REGEX.format(value))

    metadata_value = django_filters.CharFilter(
        method='filter_metadata_value', distinct=True)

    def filter_metadata_value(self, qs, name, value):
        return qs.filter(
            metadata__var_val_ucv__iregex=WORD_MATCH_REGEX.format(value))

    other_accession = django_filters.CharFilter(
        name='study__project_id',
        distinct=True)

    species = django_filters.CharFilter(method='filter_species', distinct=True)

    def filter_species(self, qs, name, value):
        return qs.filter(species__iregex=WORD_MATCH_REGEX.format(value))

    geo_loc_name = django_filters.CharFilter(
        method='filter_geo_loc_name', distinct=True)

    def filter_geo_loc_name(self, qs, name, value):
        return qs.filter(geo_loc_name__iregex=WORD_MATCH_REGEX.format(value))

    study_accession = django_filters.CharFilter(
        method='filter_study_accession', distinct=True)

    def filter_study_accession(self, qs, name, value):
        return qs.filter(
            study__accession__iregex=WORD_MATCH_REGEX.format(value))

    class Meta:
        model = emg_models.Sample
        fields = (
            'experiment_type',
            'biome',
            # 'pipeline_version',
            'geo_loc_name',
            'species',
            'instrument_model',
            'instrument_platform',
            'metadata_key',
            'metadata_value',
            'other_accession',
            'study_accession',
            'environment_material',
            'environment_feature',
        )

class RunFilter(django_filters.FilterSet):

    experiment_type = django_filters.CharFilter(
        name='experiment_type__experiment_type',
        distinct=True)

    # pipeline_version = django_filters.CharFilter(
    #     name='pipeline__release_version',
    #     distinct=True)

    biome = django_filters.CharFilter(method='filter_biome', distinct=True)

    def filter_biome(self, qs, name, value):
        try:
            b = emg_models.Biome.objects.get(lineage=value)
            qs = qs.filter(
                sample__biome__lft__gte=b.lft-1,
                sample__biome__rgt__lte=b.rgt+1)
        except emg_models.Biome.DoesNotExist:
            pass
        return qs

    species = django_filters.CharFilter(method='filter_species', distinct=True)

    def filter_species(self, qs, name, value):
        return qs.filter(
            sample__species__iregex=WORD_MATCH_REGEX.format(value))

    instrument_platform = django_filters.CharFilter(
        method='filter_instrument_platform', distinct=True)

    def filter_instrument_platform(self, qs, name, value):
        return qs.filter(
            instrument_platform__iregex=WORD_MATCH_REGEX.format(value))

    instrument_model = django_filters.CharFilter(
        method='filter_instrument_model', distinct=True)

    def filter_instrument_model(self, qs, name, value):
        return qs.filter(
            instrument_model__iregex=WORD_MATCH_REGEX.format(value))

    sample_accession = django_filters.CharFilter(
        method='filter_sample_accession', distinct=True)

    def filter_sample_accession(self, qs, name, value):
        return qs.filter(
            sample__accession__iregex=WORD_MATCH_REGEX.format(value))

    class Meta:
        model = emg_models.Run
        fields = (
            'biome',
            'experiment_type',
            # 'pipeline_version',
            'instrument_platform',
            'instrument_model',
            'species',
            'sample_accession',
        )
<commit_msg>allow filtering runs by study accession<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

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

import django_filters

from . import models as emg_models

WORD_MATCH_REGEX = r"{0}"

class PublicationFilter(django_filters.FilterSet):

    data_origination = django_filters.CharFilter(
        name='studies__data_origination',
        distinct=True)

    class Meta:
        model = emg_models.Publication
        fields = (
            'data_origination',
        )

class StudyFilter(django_filters.FilterSet):

    biome = django_filters.CharFilter(method='filter_biome', distinct=True)

    def filter_biome(self, qs, name, value):
        try:
            b = emg_models.Biome.objects.get(lineage=value)
            studies = emg_models.Sample.objects.values('study') \
                .available(self.request) \
                .filter(
                    biome__lft__gte=b.lft-1,
                    biome__rgt__lte=b.rgt+1)
            qs = qs.filter(pk__in=studies)
        except emg_models.Biome.DoesNotExist:
            pass
        return qs

    other_accession = django_filters.CharFilter(
        name='project_id',
        distinct=True)

    centre_name = django_filters.CharFilter(
        method='filter_centre_name', distinct=True)

    def filter_centre_name(self, qs, name, value):
        return qs.filter(
            centre_name__iregex=WORD_MATCH_REGEX.format(value))

    class Meta:
        model = emg_models.Study
        fields = (
            'biome',
            'centre_name',
            'other_accession',
        )

class SampleFilter(django_filters.FilterSet):

    experiment_type = django_filters.CharFilter(
        name='runs__experiment_type__experiment_type',
        distinct=True)

    # pipeline_version = django_filters.CharFilter(
    #     name='runs__pipeline__release_version',
    #     distinct=True)

    biome = django_filters.CharFilter(method='filter_biome', distinct=True)

    def filter_biome(self, qs, name, value):
        try:
            b = emg_models.Biome.objects.get(lineage=value)
            qs = qs.filter(biome__lft__gte=b.lft-1, biome__rgt__lte=b.rgt+1)
        except emg_models.Biome.DoesNotExist:
            pass
        return qs

    instrument_platform = django_filters.CharFilter(
        method='filter_instrument_platform', distinct=True)

    def filter_instrument_platform(self, qs, name, value):
        samples = emg_models.Run.objects.values('sample_id') \
            .available(self.request) \
            .filter(instrument_platform__iregex=WORD_MATCH_REGEX.format(value))
        return qs.filter(pk__in=samples).select_related('study', 'biome')

    instrument_model = django_filters.CharFilter(
        method='filter_instrument_model', distinct=True)

    def filter_instrument_model(self, qs, name, value):
        samples = emg_models.Run.objects.values('sample_id') \
            .available(self.request) \
            .filter(instrument_model__iregex=WORD_MATCH_REGEX.format(value))
        return qs.filter(pk__in=samples).select_related('study', 'biome')

    metadata_key = django_filters.CharFilter(
        method='filter_metadata_key', distinct=True)

    def filter_metadata_key(self, qs, name, value):
        return qs.filter(
            metadata__var__var_name__iregex=WORD_MATCH_REGEX.format(value))

    metadata_value = django_filters.CharFilter(
        method='filter_metadata_value', distinct=True)

    def filter_metadata_value(self, qs, name, value):
        return qs.filter(
            metadata__var_val_ucv__iregex=WORD_MATCH_REGEX.format(value))

    other_accession = django_filters.CharFilter(
        name='study__project_id',
        distinct=True)

    species = django_filters.CharFilter(method='filter_species', distinct=True)

    def filter_species(self, qs, name, value):
        return qs.filter(species__iregex=WORD_MATCH_REGEX.format(value))

    geo_loc_name = django_filters.CharFilter(
        method='filter_geo_loc_name', distinct=True)

    def filter_geo_loc_name(self, qs, name, value):
        return qs.filter(geo_loc_name__iregex=WORD_MATCH_REGEX.format(value))

    study_accession = django_filters.CharFilter(
        method='filter_study_accession', distinct=True)

    def filter_study_accession(self, qs, name, value):
        return qs.filter(
            study__accession__iregex=WORD_MATCH_REGEX.format(value))

    class Meta:
        model = emg_models.Sample
        fields = (
            'experiment_type',
            'biome',
            # 'pipeline_version',
            'geo_loc_name',
            'species',
            'instrument_model',
            'instrument_platform',
            'metadata_key',
            'metadata_value',
            'other_accession',
            'study_accession',
            'environment_material',
            'environment_feature',
        )

class RunFilter(django_filters.FilterSet):

    experiment_type = django_filters.CharFilter(
        name='experiment_type__experiment_type',
        distinct=True)

    # pipeline_version = django_filters.CharFilter(
    #     name='pipeline__release_version',
    #     distinct=True)

    biome = django_filters.CharFilter(method='filter_biome', distinct=True)

    def filter_biome(self, qs, name, value):
        try:
            b = emg_models.Biome.objects.get(lineage=value)
            qs = qs.filter(
                sample__biome__lft__gte=b.lft-1,
                sample__biome__rgt__lte=b.rgt+1)
        except emg_models.Biome.DoesNotExist:
            pass
        return qs

    species = django_filters.CharFilter(method='filter_species', distinct=True)

    def filter_species(self, qs, name, value):
        return qs.filter(
            sample__species__iregex=WORD_MATCH_REGEX.format(value))

    instrument_platform = django_filters.CharFilter(
        method='filter_instrument_platform', distinct=True)

    def filter_instrument_platform(self, qs, name, value):
        return qs.filter(
            instrument_platform__iregex=WORD_MATCH_REGEX.format(value))

    instrument_model = django_filters.CharFilter(
        method='filter_instrument_model', distinct=True)

    def filter_instrument_model(self, qs, name, value):
        return qs.filter(
            instrument_model__iregex=WORD_MATCH_REGEX.format(value))

    sample_accession = django_filters.CharFilter(
        method='filter_sample_accession', distinct=True)

    def filter_sample_accession(self, qs, name, value):
        return qs.filter(
            sample__accession__iregex=WORD_MATCH_REGEX.format(value))

    study_accession = django_filters.CharFilter(
        method='filter_study_accession', distinct=True)

    def filter_study_accession(self, qs, name, value):
        return qs.filter(
            sample__study__accession__iregex=WORD_MATCH_REGEX.format(value))

    class Meta:
        model = emg_models.Run
        fields = (
            'biome',
            'experiment_type',
            # 'pipeline_version',
            'instrument_platform',
            'instrument_model',
            'species',
            'sample_accession',
        )
<|endoftext|>


<commit_before>"""
capa freeze file format: `| capa0000 | + zlib(utf-8(json(...)))`

json format:

    {
      'version': 1,
      'base address': int(base address),
      'functions': {
        int(function va): {
          'basic blocks': {
            int(basic block va): {
              'instructions': [instruction va, ...]
            },
            ...
          },
          ...
        },
        ...
      },
      'scopes': {
        'file': [
          (str(name), [any(arg), ...], int(va), ()),
          ...
        },
        'function': [
          (str(name), [any(arg), ...], int(va), (int(function va), )),
          ...
        ],
        'basic block': [
          (str(name), [any(arg), ...], int(va), (int(function va),
                                                 int(basic block va))),
          ...
        ],
        'instruction': [
          (str(name), [any(arg), ...], int(va), (int(function va),
                                                 int(basic block va),
                                                 int(instruction va))),
          ...
        ],
      }
    }

Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
You may obtain a copy of the License at: [package root]/LICENSE.txt
Unless required by applicable law or agreed to in writing, software distributed under the License
 is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
"""
import json
import zlib
import logging

import capa.features
import capa.features.file
import capa.features.insn
import capa.features.basicblock
import capa.features.extractors
from capa.helpers import hex

logger = logging.getLogger(__name__)

def serialize_feature(feature):
    return feature.freeze_serialize()

KNOWN_FEATURES = {F.__name__: F for F in capa.features.Feature.__subclasses__()}

def deserialize_feature(doc):
    F = KNOWN_FEATURES[doc[0]]
    return F.freeze_deserialize(doc[1])

def dumps(extractor):
    """
    serialize the given extractor to a string

    args:
      extractor: capa.features.extractor.FeatureExtractor:

    returns:
      str: the serialized features.
    """
    ret = {
        "version": 1,
        "base address": extractor.get_base_address(),
        "functions": {},
        "scopes": {
            "file": [],
            "function": [],
            "basic block": [],
            "instruction": [],
        },
    }

    for feature, va in extractor.extract_file_features():
        ret["scopes"]["file"].append(serialize_feature(feature) + (hex(va), ()))

    for f in extractor.get_functions():
        ret["functions"][hex(f)] = {}

        for feature, va in extractor.extract_function_features(f):
            ret["scopes"]["function"].append(serialize_feature(feature) + (hex(va), (hex(f),)))

        for bb in extractor.get_basic_blocks(f):
            ret["functions"][hex(f)][hex(bb)] = []

            for feature, va in extractor.extract_basic_block_features(f, bb):
                ret["scopes"]["basic block"].append(
                    serialize_feature(feature)
                    + (
                        hex(va),
                        (
                            hex(f),
                            hex(bb),
                        ),
                    )
                )

            for insnva, insn in sorted(
                [(int(insn), insn) for insn in extractor.get_instructions(f, bb)], key=lambda p: p[0]
            ):
                ret["functions"][hex(f)][hex(bb)].append(hex(insnva))

                for feature, va in extractor.extract_insn_features(f, bb, insn):
                    ret["scopes"]["instruction"].append(
                        serialize_feature(feature)
                        + (
                            hex(va),
                            (
                                hex(f),
                                hex(bb),
                                hex(insnva),
                            ),
                        )
                    )
    return json.dumps(ret)

def loads(s):
    """deserialize a set of features (as a NullFeatureExtractor) from a string."""
    doc = json.loads(s)

    if doc.get("version") != 1:
        raise ValueError("unsupported freeze format version: %d" % (doc.get("version")))

    features = {
        "base address": doc.get("base address"),
        "file features": [],
        "functions": {},
    }

    for fva, function in doc.get("functions", {}).items():
        fva = int(fva, 0x10)
        features["functions"][fva] = {
            "features": [],
            "basic blocks": {},
        }

        for bbva, bb in function.items():
            bbva = int(bbva, 0x10)
            features["functions"][fva]["basic blocks"][bbva] = {
                "features": [],
                "instructions": {},
            }

            for insnva in bb:
                insnva = int(insnva, 0x10)
                features["functions"][fva]["basic blocks"][bbva]["instructions"][insnva] = {
                    "features": [],
                }

    # in the following blocks, each entry looks like:
    #
    #     ('MatchedRule', ('foo', ), '0x401000', ('0x401000', ))
    #      ^^^^^^^^^^^^^  ^^^^^^^^^  ^^^^^^^^^^  ^^^^^^^^^^^^^^
    #      feature name   args       addr         func/bb/insn
    for feature in doc.get("scopes", {}).get("file", []):
        va, loc = feature[2:]
        va = int(va, 0x10)
        feature = deserialize_feature(feature[:2])
        features["file features"].append((va, feature))

    for feature in doc.get("scopes", {}).get("function", []):
        # fetch the pair like:
        #
        #     ('0x401000', ('0x401000', ))
        #      ^^^^^^^^^^  ^^^^^^^^^^^^^^
        #      addr         func/bb/insn
        va, loc = feature[2:]
        va = int(va, 0x10)
        loc = [int(lo, 0x10) for lo in loc]

        # decode the feature from the pair like:
        #
        #     ('MatchedRule', ('foo', ))
        #      ^^^^^^^^^^^^^  ^^^^^^^^^
        #      feature name   args
        feature = deserialize_feature(feature[:2])
        features["functions"][loc[0]]["features"].append((va, feature))

    for feature in doc.get("scopes", {}).get("basic block", []):
        va, loc = feature[2:]
        va = int(va, 0x10)
        loc = [int(lo, 0x10) for lo in loc]
        feature = deserialize_feature(feature[:2])
        features["functions"][loc[0]]["basic blocks"][loc[1]]["features"].append((va, feature))

    for feature in doc.get("scopes", {}).get("instruction", []):
        va, loc = feature[2:]
        va = int(va, 0x10)
        loc = [int(lo, 0x10) for lo in loc]
        feature = deserialize_feature(feature[:2])
        features["functions"][loc[0]]["basic blocks"][loc[1]]["instructions"][loc[2]]["features"].append((va, feature))

    return capa.features.extractors.NullFeatureExtractor(features)

MAGIC = "capa0000".encode("ascii")

def dump(extractor):
    """serialize the given extractor to a byte array."""
    return MAGIC + zlib.compress(dumps(extractor).encode("utf-8"))

def is_freeze(buf):
    return buf[: len(MAGIC)] == MAGIC

def load(buf):
    """deserialize a set of features (as a NullFeatureExtractor) from a byte array."""
    if not is_freeze(buf):
        raise ValueError("missing magic header")
    return loads(zlib.decompress(buf[len(MAGIC) :]).decode("utf-8"))

def main(argv=None):
    import sys
    import argparse

    import capa.main

    if argv is None:
        argv = sys.argv[1:]

    formats = [
        ("auto", "(default) detect file type automatically"),
        ("pe", "Windows PE file"),
        ("sc32", "32-bit shellcode"),
        ("sc64", "64-bit shellcode"),
    ]
    format_help = ", ".join(["%s: %s" % (f[0], f[1]) for f in formats])

    parser = argparse.ArgumentParser(description="save capa features to a file")
    parser.add_argument("sample", type=str, help="Path to sample to analyze")
    parser.add_argument("output", type=str, help="Path to output file")
    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output")
    parser.add_argument("-q", "--quiet", action="store_true", help="Disable all output but errors")
    parser.add_argument(
        "-f", "--format", choices=[f[0] for f in formats], default="auto", help="Select sample format, %s" % format_help
    )
    parser.add_argument(
        "-b",
        "--backend",
        type=str,
        help="select the backend to use",
        choices=(capa.main.BACKEND_VIV, capa.main.BACKEND_SMDA),
        default=capa.main.BACKEND_VIV,
    )
    parser.add_argument(
        "--signature",
        action="append",
        dest="signatures",
        type=str,
        default=[],
        help="use the given signatures to identify library functions, file system paths to .sig/.pat files.",
    )
    args = parser.parse_args(args=argv)

    if args.quiet:
        logging.basicConfig(level=logging.ERROR)
        logging.getLogger().setLevel(logging.ERROR)
    elif args.verbose:
        logging.basicConfig(level=logging.DEBUG)
        logging.getLogger().setLevel(logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)
        logging.getLogger().setLevel(logging.INFO)

    extractor = capa.main.get_extractor(args.sample, args.format, args.backend, sigpaths=args.signatures)
    with open(args.output, "wb") as f:
        f.write(dump(extractor))

    return 0

if __name__ == "__main__":
    import sys

    sys.exit(main())
<commit_msg>freeze: use common args<commit_after>"""
capa freeze file format: `| capa0000 | + zlib(utf-8(json(...)))`

json format:

    {
      'version': 1,
      'base address': int(base address),
      'functions': {
        int(function va): {
          'basic blocks': {
            int(basic block va): {
              'instructions': [instruction va, ...]
            },
            ...
          },
          ...
        },
        ...
      },
      'scopes': {
        'file': [
          (str(name), [any(arg), ...], int(va), ()),
          ...
        },
        'function': [
          (str(name), [any(arg), ...], int(va), (int(function va), )),
          ...
        ],
        'basic block': [
          (str(name), [any(arg), ...], int(va), (int(function va),
                                                 int(basic block va))),
          ...
        ],
        'instruction': [
          (str(name), [any(arg), ...], int(va), (int(function va),
                                                 int(basic block va),
                                                 int(instruction va))),
          ...
        ],
      }
    }

Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
You may obtain a copy of the License at: [package root]/LICENSE.txt
Unless required by applicable law or agreed to in writing, software distributed under the License
 is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
"""
import json
import zlib
import logging

import capa.features
import capa.features.file
import capa.features.insn
import capa.features.basicblock
import capa.features.extractors
from capa.helpers import hex

logger = logging.getLogger(__name__)

def serialize_feature(feature):
    return feature.freeze_serialize()

KNOWN_FEATURES = {F.__name__: F for F in capa.features.Feature.__subclasses__()}

def deserialize_feature(doc):
    F = KNOWN_FEATURES[doc[0]]
    return F.freeze_deserialize(doc[1])

def dumps(extractor):
    """
    serialize the given extractor to a string

    args:
      extractor: capa.features.extractor.FeatureExtractor:

    returns:
      str: the serialized features.
    """
    ret = {
        "version": 1,
        "base address": extractor.get_base_address(),
        "functions": {},
        "scopes": {
            "file": [],
            "function": [],
            "basic block": [],
            "instruction": [],
        },
    }

    for feature, va in extractor.extract_file_features():
        ret["scopes"]["file"].append(serialize_feature(feature) + (hex(va), ()))

    for f in extractor.get_functions():
        ret["functions"][hex(f)] = {}

        for feature, va in extractor.extract_function_features(f):
            ret["scopes"]["function"].append(serialize_feature(feature) + (hex(va), (hex(f),)))

        for bb in extractor.get_basic_blocks(f):
            ret["functions"][hex(f)][hex(bb)] = []

            for feature, va in extractor.extract_basic_block_features(f, bb):
                ret["scopes"]["basic block"].append(
                    serialize_feature(feature)
                    + (
                        hex(va),
                        (
                            hex(f),
                            hex(bb),
                        ),
                    )
                )

            for insnva, insn in sorted(
                [(int(insn), insn) for insn in extractor.get_instructions(f, bb)], key=lambda p: p[0]
            ):
                ret["functions"][hex(f)][hex(bb)].append(hex(insnva))

                for feature, va in extractor.extract_insn_features(f, bb, insn):
                    ret["scopes"]["instruction"].append(
                        serialize_feature(feature)
                        + (
                            hex(va),
                            (
                                hex(f),
                                hex(bb),
                                hex(insnva),
                            ),
                        )
                    )
    return json.dumps(ret)

def loads(s):
    """deserialize a set of features (as a NullFeatureExtractor) from a string."""
    doc = json.loads(s)

    if doc.get("version") != 1:
        raise ValueError("unsupported freeze format version: %d" % (doc.get("version")))

    features = {
        "base address": doc.get("base address"),
        "file features": [],
        "functions": {},
    }

    for fva, function in doc.get("functions", {}).items():
        fva = int(fva, 0x10)
        features["functions"][fva] = {
            "features": [],
            "basic blocks": {},
        }

        for bbva, bb in function.items():
            bbva = int(bbva, 0x10)
            features["functions"][fva]["basic blocks"][bbva] = {
                "features": [],
                "instructions": {},
            }

            for insnva in bb:
                insnva = int(insnva, 0x10)
                features["functions"][fva]["basic blocks"][bbva]["instructions"][insnva] = {
                    "features": [],
                }

    # in the following blocks, each entry looks like:
    #
    #     ('MatchedRule', ('foo', ), '0x401000', ('0x401000', ))
    #      ^^^^^^^^^^^^^  ^^^^^^^^^  ^^^^^^^^^^  ^^^^^^^^^^^^^^
    #      feature name   args       addr         func/bb/insn
    for feature in doc.get("scopes", {}).get("file", []):
        va, loc = feature[2:]
        va = int(va, 0x10)
        feature = deserialize_feature(feature[:2])
        features["file features"].append((va, feature))

    for feature in doc.get("scopes", {}).get("function", []):
        # fetch the pair like:
        #
        #     ('0x401000', ('0x401000', ))
        #      ^^^^^^^^^^  ^^^^^^^^^^^^^^
        #      addr         func/bb/insn
        va, loc = feature[2:]
        va = int(va, 0x10)
        loc = [int(lo, 0x10) for lo in loc]

        # decode the feature from the pair like:
        #
        #     ('MatchedRule', ('foo', ))
        #      ^^^^^^^^^^^^^  ^^^^^^^^^
        #      feature name   args
        feature = deserialize_feature(feature[:2])
        features["functions"][loc[0]]["features"].append((va, feature))

    for feature in doc.get("scopes", {}).get("basic block", []):
        va, loc = feature[2:]
        va = int(va, 0x10)
        loc = [int(lo, 0x10) for lo in loc]
        feature = deserialize_feature(feature[:2])
        features["functions"][loc[0]]["basic blocks"][loc[1]]["features"].append((va, feature))

    for feature in doc.get("scopes", {}).get("instruction", []):
        va, loc = feature[2:]
        va = int(va, 0x10)
        loc = [int(lo, 0x10) for lo in loc]
        feature = deserialize_feature(feature[:2])
        features["functions"][loc[0]]["basic blocks"][loc[1]]["instructions"][loc[2]]["features"].append((va, feature))

    return capa.features.extractors.NullFeatureExtractor(features)

MAGIC = "capa0000".encode("ascii")

def dump(extractor):
    """serialize the given extractor to a byte array."""
    return MAGIC + zlib.compress(dumps(extractor).encode("utf-8"))

def is_freeze(buf):
    return buf[: len(MAGIC)] == MAGIC

def load(buf):
    """deserialize a set of features (as a NullFeatureExtractor) from a byte array."""
    if not is_freeze(buf):
        raise ValueError("missing magic header")
    return loads(zlib.decompress(buf[len(MAGIC) :]).decode("utf-8"))

def main(argv=None):
    import sys
    import argparse

    import capa.main

    if argv is None:
        argv = sys.argv[1:]

    parser = argparse.ArgumentParser(description="save capa features to a file")
    capa.main.install_common_args(parser, {"sample", "format", "backend", "signatures"})
    parser.add_argument("output", type=str, help="Path to output file")
    args = parser.parse_args(args=argv)
    capa.main.handle_common_args(args)

    extractor = capa.main.get_extractor(args.sample, args.format, args.backend, sigpaths=args.signatures)
    with open(args.output, "wb") as f:
        f.write(dump(extractor))

    return 0

if __name__ == "__main__":
    import sys

    sys.exit(main())
<|endoftext|>