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

#This software is distributed under the Creative Commons license (CC0) version 1.0. A copy of this license should have been distributed with this software.
#The license can also be read online: <https://creativecommons.org/publicdomain/zero/1.0/>. If this online license differs from the license provided with this software, the license provided with this software should be applied.

"""
Tests the listening module that provides a way to listen for state changes.
"""

import unittest #To define automatic tests.
import unittest.mock #To track how often a listener function was called.

import luna.listen #The module we're testing.

class TestListen(unittest.TestCase):
	"""
	Tests the listening module that provides a way to listen for state changes.
	"""

	def setUp(self):
		"""
		Prepares a few fields on this object of which we can track changes.
		"""
		self.field_integer = 0
		self.field_string = ""

	def tearDown(self):
		"""
		Cleans up the object after possibly modifying it.
		"""
		if hasattr(self, "field_float"):
			delattr(self, "field_float")

	def test_listen_all_fields(self):
		"""
		Tests listening to all changes of an instance.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self)
		self.field_integer = 1
		listener.assert_called_once_with("field_integer", 1)
		self.field_string = "I love you."
		listener.assert_called_with("field_string", "I love you.")

	def test_listen_new_fields(self):
		"""
		Test listening to all changes of an instance and then making new fields.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self)
		self.field_float = 3.1416 #pylint: disable=attribute-defined-outside-init
		listener.assert_called_once_with("field_float", 3.1416)

	def test_listen_nochange(self):
		"""
		Tests that the listener doesn't get called if the state doesn't change.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self, "field_integer")
		self.field_integer = 0 #Equal to starting state.
		listener.assert_not_called()

	def test_listen_simple(self):
		"""
		Tests a simple happy path with a single state change and single
		listener.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self, "field_integer")
		self.field_integer = 1 #Triggers a change.
		listener.assert_called_once_with("field_integer", 1)

	def test_listen_twice(self):
		"""
		Tests listening for two consecutive state changes.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self, "field_integer")
		self.field_integer = 1
		listener.assert_called_once_with("field_integer", 1)
		self.field_integer = 0 #Trigger two changes.
		listener.assert_called_with("field_integer", 0)
		self.assertEqual(listener.call_count, 2, "The state was changed twice.")<commit_msg>Add test for listening to multiple attributes at the same time<commit_after>#!/usr/bin/env python
#-*- coding: utf-8 -*-

#This software is distributed under the Creative Commons license (CC0) version 1.0. A copy of this license should have been distributed with this software.
#The license can also be read online: <https://creativecommons.org/publicdomain/zero/1.0/>. If this online license differs from the license provided with this software, the license provided with this software should be applied.

"""
Tests the listening module that provides a way to listen for state changes.
"""

import unittest #To define automatic tests.
import unittest.mock #To track how often a listener function was called.

import luna.listen #The module we're testing.

class TestListen(unittest.TestCase):
	"""
	Tests the listening module that provides a way to listen for state changes.
	"""

	def setUp(self):
		"""
		Prepares a few fields on this object of which we can track changes.
		"""
		self.field_integer = 0
		self.field_string = ""

	def tearDown(self):
		"""
		Cleans up the object after possibly modifying it.
		"""
		if hasattr(self, "field_float"):
			delattr(self, "field_float")

	def test_listen_all_fields(self):
		"""
		Tests listening to all changes of an instance.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self)
		self.field_integer = 1
		listener.assert_called_once_with("field_integer", 1)
		self.field_string = "I love you."
		listener.assert_called_with("field_string", "I love you.")

	def test_listen_new_fields(self):
		"""
		Test listening to all changes of an instance and then making new fields.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self)
		self.field_float = 3.1416 #pylint: disable=attribute-defined-outside-init
		listener.assert_called_once_with("field_float", 3.1416)

	def test_listen_nochange(self):
		"""
		Tests that the listener doesn't get called if the state doesn't change.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self, "field_integer")
		self.field_integer = 0 #Equal to starting state.
		listener.assert_not_called()

	def test_listen_simple(self):
		"""
		Tests a simple happy path with a single state change and single
		listener.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self, "field_integer")
		self.field_integer = 1 #Triggers a change.
		listener.assert_called_once_with("field_integer", 1)

	def test_listen_twice(self):
		"""
		Tests listening for two consecutive state changes.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self, "field_integer")
		self.field_integer = 1
		listener.assert_called_once_with("field_integer", 1)
		self.field_integer = 0 #Trigger two changes.
		listener.assert_called_with("field_integer", 0)
		self.assertEqual(listener.call_count, 2, "The state was changed twice.")

	def test_listen_multiple_attributes(self):
		"""
		Tests listening for two attributes and the instance at the same time.
		"""
		listener = unittest.mock.MagicMock()
		luna.listen.listen(listener, self, "field_integer")
		luna.listen.listen(listener, self, "field_string")
		luna.listen.listen(listener, self)
		self.field_integer = 1
		listener.assert_called_with("field_integer", 1)
		self.assertEqual(listener.call_count, 2, "The listener must be called once for the attribute and once for the instance.")
		self.field_string = "poo"
		listener.assert_called_with("field_string", "poo")
		self.assertEqual(listener.call_count, 4, "The listener must be called twice for two attributes and twice for the instance.")<|endoftext|>


<commit_before>from nipype.interfaces import fsl
from nipype.algorithms.modelgen import  SpecifyModel
from nipype.interfaces.utility import IdentityInterface, Function, Rename
from nipype.pipeline.engine import Node, MapNode, Workflow

def create_timeseries_model_workflow(name="model", exp_info={}):

    inputnode = Node(IdentityInterface(fields=["subject_id",
                                               "outlier_files",
                                               "mean_func",
                                               "realign_params",
                                               "timeseries"]),
                     name="inputnode")

    # Read in par files to determine model information
    modelinfo = Node(Function(input_names=["subject_id",
                                           "functional_runs",
                                           "exp_info"],
                              output_names=["subject_info"],
                              function=build_model_info),
                     name="modelinfo")
    modelinfo.inputs.exp_info = exp_info

    # Generate Nipype-style model specification
    modelspec = Node(SpecifyModel(
        time_repetition=exp_info["TR"],
        input_units=exp_info["units"],
        high_pass_filter_cutoff=exp_info["hpf_cutoff"]),
                     name="modelspec")

    # Generate FSL fsf model specifcations
    level1design = Node(fsl.Level1Design(model_serial_correlations=True,
                                         bases=exp_info["hrf_bases"],
                                         interscan_interval=exp_info["TR"],
                                         contrasts=exp_info["contrasts"]),
                        name="level1design")

    # Generate design and contrast matrix files
    featmodel = MapNode(fsl.FEATModel(),
                        iterfield=["fsf_file", "ev_files"],
                        name="featmodel")

    # Generate a plot of regressor correlation
    designcorr = MapNode(Function(input_names=["in_file",
                                               "ev_files",
                                               "n_runs"],
                                  output_names=["out_file"],
                                  function=design_corr),
                         iterfield=["in_file"],
                         name="designcorr")
    designcorr.inputs.n_runs = exp_info["n_runs"]

    # Rename the design matrix, contrasts, and image
    rename_design = MapNode(Rename(format_string="design",
                                   keep_ext=True),
                            iterfield=["in_file"],
                            name="rename_design")

    rename_x_mat = MapNode(Rename(format_string="design",
                                  keep_ext=True),
                           iterfield=["in_file"],
                           name="rename_x_mat")

    rename_c_mat = MapNode(Rename(format_string="design",
                                  keep_ext=True),
                           iterfield=["in_file"],
                           name="rename_c_mat")

    # Use film_gls to estimate the timeseries model
    modelestimate = MapNode(fsl.FILMGLS(smooth_autocorr=True,
                                        mask_size=5,
                                        threshold=1000),
                            iterfield=["design_file", "in_file"],
                            name="modelestimate")

    # Plot the residual error map for quality control
    plotresidual = MapNode(Function(input_names=["resid_file",
                                                 "background_file"],
                                   output_names=["out_file"],
                                   function=plot_residual),
                          iterfield=["resid_file", "background_file"],
                          name="plotresidual")

    # Run the contrast estimation routine
    contrastestimate = MapNode(fsl.ContrastMgr(),
                               iterfield=["tcon_file",
                                          "dof_file",
                                          "corrections",
                                          "param_estimates",
                                          "sigmasquareds"],
                               name="contrastestimate")

    # Plot the zstat images
    plotzstats = MapNode(Function(input_names=["background_file",
                                               "zstat_files",
                                               "contrasts"],
                                  output_names=["out_files"],
                                  function=plot_zstats),
                         iterfield=["background_file", "zstat_files"],
                         name="plotzstats")
    plotzstats.inputs.contrasts = exp_info["contrast_names"]

    # Build pdf and html reports
    report = MapNode(Function(input_names=["subject_id",
                                           "design_image",
                                           "design_corr",
                                           "residual",
                                           "zstat_pngs",
                                           "contrast_names"],
                              output_names=["reports"],
                              function=write_model_report),
                     iterfield=["design_image",
                                "design_corr",
                                "residual",
                                "zstat_pngs"],
                     name="report")
    report.inputs.contrast_names = exp_info["contrast_names"]

    # Dump the exp_info dictionary to json
    dumpjson = MapNode(Function(input_names=["exp_info",
                                             "timeseries"],
                                output_names=["json_file"],
                                function=dump_exp_info),
                       iterfield=["timeseries"],
                       name="dumpjson")
    dumpjson.inputs.exp_info = exp_info

    # Define the workflow outputs
    outputnode = Node(IdentityInterface(fields=["results",
                                                "design_image",
                                                "design_corr",
                                                "sigmasquareds",
                                                "copes",
                                                "varcopes",
                                                "zstats",
                                                "reports",
                                                "design_mat",
                                                "contrast_mat",
                                                "json_file",
                                                "zstat_pngs"]),
                      name="outputnode")

    # Define the workflow and connect the nodes
    model = Workflow(name=name)
    model.connect([
        (inputnode, modelinfo,
            [("subject_id", "subject_id"),
             ("timeseries", "functional_runs")]),
        (modelinfo, modelspec,
            [("subject_info", "subject_info")]),
        (inputnode, modelspec,
            [("timeseries", "functional_runs"),
             ("outlier_files", "outlier_files"),
             ("realign_params", "realignment_parameters")]),
        (inputnode, modelestimate,
            [("timeseries", "in_file")]),
        (modelspec, level1design,
            [("session_info", "session_info")]),
        (level1design, featmodel,
            [("fsf_files", "fsf_file"),
             ("ev_files", "ev_files")]),
        (featmodel, modelestimate,
            [("design_file", "design_file")]),
        (featmodel, contrastestimate,
            [("con_file", "tcon_file")]),
        (featmodel, designcorr,
            [("design_file", "in_file")]),
        (featmodel, rename_x_mat,
            [("design_file", "in_file")]),
        (featmodel, rename_c_mat,
            [("con_file", "in_file")]),
        (level1design, designcorr,
            [("ev_files", "ev_files")]),
        (featmodel, rename_design,
            [("design_image", "in_file")]),
        (modelestimate, plotresidual,
            [("sigmasquareds", "resid_file")]),
        (inputnode, plotresidual,
            [("mean_func", "background_file")]),
        (modelestimate, contrastestimate,
            [("dof_file", "dof_file"),
             ("corrections", "corrections"),
             ("param_estimates", "param_estimates"),
             ("sigmasquareds", "sigmasquareds")]),
        (contrastestimate, plotzstats,
            [("zstats", "zstat_files")]),
        (inputnode, plotzstats,
            [("mean_func", "background_file")]),
        (inputnode, dumpjson,
            [("timeseries", "timeseries")]),
        (inputnode, report,
            [("subject_id", "subject_id")]),
        (rename_design, report,
            [("out_file", "design_image")]),
        (designcorr, report,
            [("out_file", "design_corr")]),
        (plotresidual, report,
            [("out_file", "residual")]),
        (plotzstats, report,
            [("out_files", "zstat_pngs")]),
        (report, outputnode,
            [("reports", "reports")]),
        (dumpjson, outputnode,
            [("json_file", "json_file")]),
        (rename_design, outputnode,
            [("out_file", "design_image")]),
        (rename_x_mat, outputnode,
            [("out_file", "design_mat")]),
        (rename_c_mat, outputnode,
            [("out_file", "contrast_mat")]),
        (designcorr, outputnode,
            [("out_file", "design_corr")]),
        (plotresidual, outputnode,
            [("out_file", "sigmasquareds")]),
        (modelestimate, outputnode,
            [("results_dir", "results")]),
        (contrastestimate, outputnode,
            [("copes", "copes"),
             ("varcopes", "varcopes"),
             ("zstats", "zstats")]),
        (plotzstats, outputnode,
            [("out_files", "zstat_pngs")]),
        ])

    return model, inputnode, outputnode

def build_model_info(subject_id, functional_runs, exp_info):
    import os.path as op
    from copy import deepcopy
    from numpy import loadtxt, atleast_1d, atleast_2d
    from nipype.interfaces.base import Bunch

    events = exp_info["events"]
    regressor_names = exp_info["regressors"]

    model_info = []

    if not isinstance(functional_runs, list):
        functional_runs = [functional_runs]
    n_runs = len(functional_runs)
    for run in range(1, n_runs + 1):
        onsets, durations, amplitudes, regressors = [], [], [], []
        for event in events:
            event_info = dict(event=event, run=run, subject_id=subject_id)
            parfile = op.join(exp_info["parfile_base_dir"],
                              exp_info['parfile_template'] % event_info)
            o, d, a = atleast_2d(loadtxt(parfile)).T
            onsets.append(o)
            durations.append(d)
            amplitudes.append(a)
        for regressor in regressor_names:
            regress_info = dict(regressor=regressor,
                                run=run,
                                subject_id=subject_id)
            regressor_file = op.join(
                exp_info["regressor_base_dir"],
                exp_info["regressor_template"] % regress_info)
            regressors.append(atleast_1d(loadtxt(regressor_file)))

        model_info.append(
            Bunch(conditions=events,
                  regressor_names=regressor_names,
                  regressors=deepcopy(regressors),
                  onsets=deepcopy(onsets),
                  durations=deepcopy(durations),
                  amplitudes=deepcopy(amplitudes)))

    return model_info

def design_corr(in_file, ev_files, n_runs):
    import re
    from os.path import abspath, basename
    import numpy as np
    import matplotlib.pyplot as plt
    X = np.loadtxt(in_file, skiprows=5)
    f = plt.figure(figsize=(5, 5))
    ax = f.add_subplot(111)
    ax.matshow(np.abs(np.corrcoef(X.T)), vmin=0, vmax=1, cmap="hot")
    run = int(re.match(r"run(\d+).mat", basename(in_file)).group(1))
    pat = "ev_(\w+)_%d_\d+.txt" % run
    if n_runs > 1:
        ev_files = ev_files[run]
    ev_names = [re.match(pat, basename(f)).group(1) for f in ev_files]
    ev_names = map(lambda x: x.lower(), ev_names)
    ax.set_xticks(range(len(ev_names)))
    ax.set_yticks(range(len(ev_names)))
    ax.set_xticklabels(ev_names, fontsize=6, rotation="vertical")
    ax.set_yticklabels(ev_names, fontsize=6)
    out_file = abspath("design_correlation.png")
    plt.savefig(out_file)
    return out_file

def plot_residual(resid_file, background_file):
    """Plot the model residual. Use FSL for now; should move to pylab."""
    from os.path import abspath
    from nibabel import load
    from subprocess import call
    from scipy.stats import scoreatpercentile
    ov_nii = "sigmasquareds_overlay.nii.gz"
    out_file = abspath("sigmasquareds.png")
    err_img = load(resid_file)
    err_data = err_img.get_data()
    low = scoreatpercentile(err_data.ravel(), 2)
    high = scoreatpercentile(err_data.ravel(), 98)
    n_slice = err_img.shape[-1]
    native = n_slice < 50  # arbitrary but should work fine
    width = 750 if native else 872
    sample = 1 if native else 2
    del err_data
    call(["overlay", "1", "0", background_file, "-a",
          resid_file, str(low), str(high), ov_nii])
    call(["slicer", ov_nii, "-S", str(sample), str(width), out_file])
    return out_file

def plot_zstats(background_file, zstat_files, contrasts):
    """Plot the zstats and return a directory of pngs."""
    from os.path import abspath
    from nibabel import load
    from subprocess import call
    bg_img = load(background_file)
    n_slice = bg_img.shape[-1]
    native = n_slice < 50
    width = 750 if native else 872
    sample = 1 if native else 2
    out_files = []
    if not isinstance(zstat_files, list):
        zstat_files = [zstat_files]
    for i, contrast in enumerate(contrasts):
        zstat_file = zstat_files[i]
        zstat_png = "zstat%d.png" % (i + 1)
        ov_nii = "zstat%d_overlay.nii.gz" % (i + 1)
        call(["overlay", "1", "0", background_file, "-a",
              zstat_file, "2.3", "10",
              zstat_file, "-2.3", "-10", ov_nii])
        call(["slicer", ov_nii, "-S", str(sample), str(width), zstat_png])
        out_files.append(abspath(zstat_png))
    return out_files

def dump_exp_info(exp_info, timeseries):
    """Dump the exp_info dict into a json file."""
    from copy import deepcopy
    from os.path import abspath
    import json
    json_file = abspath("experiment_info.json")
    json_info = deepcopy(exp_info)
    del json_info["contrasts"]
    with open(json_file, "w") as fp:
        json.dump(json_info, fp, sort_keys=True, indent=4)
    return json_file

def write_model_report(subject_id, design_image, design_corr,
                       residual, zstat_pngs, contrast_names):
    """Write model report info to rst and convert to pdf/html."""
    import time
    from lyman.tools import write_workflow_report
    from lyman.workflows.reporting import model_report_template

    # Fill in the initial report template dict
    report_dict = dict(now=time.asctime(),
                       subject_id=subject_id,
                       design_image=design_image,
                       design_corr=design_corr,
                       residual=residual)

    # Add the zstat image templates and update the dict
    for i, con in enumerate(contrast_names, 1):
        report_dict["con%d_name" % i] = con
        report_dict["zstat%d_png" % i] = zstat_pngs[i - 1]
        header = "Zstat %d: %s" % (i, con)
        model_report_template = "\n".join(
            [model_report_template,
             header,
             "".join(["^" for l in header]),
             "",
             "".join([".. image:: %(zstat", str(i), "_png)s"]),
             "    :width: 6.5in",
             ""])

    out_files = write_workflow_report("model",
                                      model_report_template,
                                      report_dict)
    return out_files
<commit_msg>Starting model refactor<commit_after>"""Timeseries model using FSL's gaussian least squares."""
import os
import os.path as op
import re
import numpy as np
import scipy as sp
from scipy import stats, signal
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import moss
from moss import glm
import seaborn as sns

from nipype import fsl, Node, MapNode, Workflow, IdentityInterface, Function

from lyman import default_experiment_parameters

imports = ["import os",
           "import os.path as op",
           "import re",
           "import numpy as np",
           "import scipy as sp",
           "from scipy import stats, signal",
           "import pandas as pd",
           "import matplotlib as mpl",
           "import matplotlib.pyplot as plt",
           "import moss",
           "from moss import glm",
           "import seaborn as sns"]

def create_timeseries_model_workflow(name="model", exp_info=None):

    # Node inputs
    inputnode = Node(IdentityInterface(["subject_id",
                                        "design_file",
                                        "realign_file",
                                        "artifact_file",
                                        "mean_file",
                                        "timeseries"]),
                     "inputs")

    # Default experiment parameters for generating graph inamge, testing, etc.
    if exp_info is None:
        exp_info = default_experiment_parameters()

    # Set up the experimental design
    modeldesign = MapNode(Function(["exp_info",
                                    "design_file",
                                    "realign_file",
                                    "artifact_file",
                                    "run"],
                                   ["design_matrix_file",
                                    "contrast_file",
                                    "report"],
                                   design_model,
                                   imports),
                          [],
                          "modeldesign")

    # Use film_gls to estimate the timeseries model
    modelestimate = MapNode(fsl.FILMGLS(smooth_autocorr=True,
                                        mask_size=5,
                                        threshold=1000),
                            ["design_file", "in_file"],
                            "modelestimate")

    # Run the contrast estimation routine
    contrastestimate = MapNode(fsl.ContrastMgr(),
                               ["tcon_file",
                                "dof_file",
                                "corrections",
                                "param_estimates",
                                "sigmasquareds"],
                               "contrastestimate")

    # Define the workflow outputs
    outputnode = Node(IdentityInterface(["results",
                                         "design_image",
                                         "design_corr",
                                         "sigmasquareds",
                                         "copes",
                                         "varcopes",
                                         "zstats",
                                         "reports",
                                         "design_mat",
                                         "contrast_mat",
                                         "json_file",
                                         "zstat_pngs"]),
                      "outputs")

    # Define the workflow and connect the nodes
    model = Workflow(name=name)
    model.connect([
        (inputnode, modelinfo,
            [("subject_id", "subject_id"),
             ("timeseries", "functional_runs")]),
        (modelinfo, modelspec,
            [("subject_info", "subject_info")]),
        (inputnode, modelspec,
            [("timeseries", "functional_runs"),
             ("outlier_files", "outlier_files"),
             ("realign_params", "realignment_parameters")]),
        (inputnode, modelestimate,
            [("timeseries", "in_file")]),
        (modelspec, level1design,
            [("session_info", "session_info")]),
        (level1design, featmodel,
            [("fsf_files", "fsf_file"),
             ("ev_files", "ev_files")]),
        (featmodel, modelestimate,
            [("design_file", "design_file")]),
        (featmodel, contrastestimate,
            [("con_file", "tcon_file")]),
        (featmodel, designcorr,
            [("design_file", "in_file")]),
        (featmodel, rename_x_mat,
            [("design_file", "in_file")]),
        (featmodel, rename_c_mat,
            [("con_file", "in_file")]),
        (level1design, designcorr,
            [("ev_files", "ev_files")]),
        (featmodel, rename_design,
            [("design_image", "in_file")]),
        (modelestimate, plotresidual,
            [("sigmasquareds", "resid_file")]),
        (inputnode, plotresidual,
            [("mean_func", "background_file")]),
        (modelestimate, contrastestimate,
            [("dof_file", "dof_file"),
             ("corrections", "corrections"),
             ("param_estimates", "param_estimates"),
             ("sigmasquareds", "sigmasquareds")]),
        (contrastestimate, plotzstats,
            [("zstats", "zstat_files")]),
        (inputnode, plotzstats,
            [("mean_func", "background_file")]),
        (inputnode, dumpjson,
            [("timeseries", "timeseries")]),
        (inputnode, report,
            [("subject_id", "subject_id")]),
        (rename_design, report,
            [("out_file", "design_image")]),
        (designcorr, report,
            [("out_file", "design_corr")]),
        (plotresidual, report,
            [("out_file", "residual")]),
        (plotzstats, report,
            [("out_files", "zstat_pngs")]),
        (report, outputnode,
            [("reports", "reports")]),
        (dumpjson, outputnode,
            [("json_file", "json_file")]),
        (rename_design, outputnode,
            [("out_file", "design_image")]),
        (rename_x_mat, outputnode,
            [("out_file", "design_mat")]),
        (rename_c_mat, outputnode,
            [("out_file", "contrast_mat")]),
        (designcorr, outputnode,
            [("out_file", "design_corr")]),
        (plotresidual, outputnode,
            [("out_file", "sigmasquareds")]),
        (modelestimate, outputnode,
            [("results_dir", "results")]),
        (contrastestimate, outputnode,
            [("copes", "copes"),
             ("varcopes", "varcopes"),
             ("zstats", "zstats")]),
        (plotzstats, outputnode,
            [("out_files", "zstat_pngs")]),
        ])

    return model, inputnode, outputnode

# Main Interface Functions
# ========================

def design_model(design_file, realign_file, artifact_file, exp_info, run):
    """Build the model design."""
    design = pd.read_csv(design_file)
    design = design[design.run == run]

    realign = pd.read_csv(realign_file)
    realign = realign.filter(regex="rot|trans").apply(stats.zscore)

    artifacts = pd.read_csv(artifact_file).max(axis=1)
    ntp = len(artifacts)
    tr = exp_info["TR"]

    # Set up the HRF model
    hrf = getattr(glm, exp_info["hrf_model"])
    hrf = hrf(exp_info["temporal_deriv"], tr, **exp_info["hrf_params"])

    # Keep tabs on the keyword arguments for the design matrix
    design_kwargs = dict(confounds=realign,
                         artifacts=artifacts,
                         tr=tr,
                         condition_names=exp_info["condition_names"],
                         confound_pca=exp_info["confound_pca"],
                         hpf_cutoff=exp_info["hpf_cutoff"])

    # Create the main design matrix object
    X = glm.DesignMatrix(design, hrf, ntp, **design_kwargs)

    # Also create one without high pass filtering
    design_kwargs["hpf_cutoff"] = None
    X_unfiltered = glm.DesignMatrix(design, hrf, ntp, **design_kwargs)

    # Now we build up the report, with lots of lovely images
    sns.set()
    design_png = op.abspath("design.png")
    X.plot(fname=design_png)

    corr_png = op.abspath("design_correlation.png")
    X.plot_confound_correlation(fname=corr_png)

    svd_png = op.abspath("design_singular_values.png")
    X.plot_singular_values(fname=svd_png)

    report = [design_png, corr_png, svd_png]

    # Now plot the information loss from the filter
    for i, (name, cols, weights) in enumerate(exp_info["contrasts"], 1):
        C = X.contrast_vector(cols, weights)

        y_filt = X.design_matrix.dot(C)
        y_unfilt = X_unfiltered.design_matrix.dot(C)

        fs, pxx_filt = signal.welch(y_filt, 1. / tr, nperseg=ntp)
        fs, pxx_unfilt = signal.welch(y_unfilt, 1. / tr, nperseg=ntp)

        f, ax = plt.subplots(1, 1)
        ax.fill_between(fs, pxx_unfilt, color="#C41E3A")
        ax.fill_between(fs, pxx_filt, color="#444444")
        ax.set_xlabel("Frequency")
        ax.set_ylabel("Spectral Density")
        plt.tight_layout()
        fname = op.abspath("cope%d_filter.png" % i)
        f.savefig(fname)
        report.append(fname)

    # Finally, write out the design files in FSL format
    design_matrix_file = op.abspath("design.mat")
    contrast_file = op.abspath("design.con")
    X.to_fsl_files("design", exp_info["contrasts"])

    return design_matrix_file, contrast_file, report

def build_model_info(subject_id, functional_runs, exp_info):
    import os.path as op
    from copy import deepcopy
    from numpy import loadtxt, atleast_1d, atleast_2d
    from nipype.interfaces.base import Bunch

    events = exp_info["events"]
    regressor_names = exp_info["regressors"]

    model_info = []

    if not isinstance(functional_runs, list):
        functional_runs = [functional_runs]
    n_runs = len(functional_runs)
    for run in range(1, n_runs + 1):
        onsets, durations, amplitudes, regressors = [], [], [], []
        for event in events:
            event_info = dict(event=event, run=run, subject_id=subject_id)
            parfile = op.join(exp_info["parfile_base_dir"],
                              exp_info['parfile_template'] % event_info)
            o, d, a = atleast_2d(loadtxt(parfile)).T
            onsets.append(o)
            durations.append(d)
            amplitudes.append(a)
        for regressor in regressor_names:
            regress_info = dict(regressor=regressor,
                                run=run,
                                subject_id=subject_id)
            regressor_file = op.join(
                exp_info["regressor_base_dir"],
                exp_info["regressor_template"] % regress_info)
            regressors.append(atleast_1d(loadtxt(regressor_file)))

        model_info.append(
            Bunch(conditions=events,
                  regressor_names=regressor_names,
                  regressors=deepcopy(regressors),
                  onsets=deepcopy(onsets),
                  durations=deepcopy(durations),
                  amplitudes=deepcopy(amplitudes)))

    return model_info

def design_corr(in_file, ev_files, n_runs):
    import re
    from os.path import abspath, basename
    import numpy as np
    import matplotlib.pyplot as plt
    X = np.loadtxt(in_file, skiprows=5)
    f = plt.figure(figsize=(5, 5))
    ax = f.add_subplot(111)
    ax.matshow(np.abs(np.corrcoef(X.T)), vmin=0, vmax=1, cmap="hot")
    run = int(re.match(r"run(\d+).mat", basename(in_file)).group(1))
    pat = "ev_(\w+)_%d_\d+.txt" % run
    if n_runs > 1:
        ev_files = ev_files[run]
    ev_names = [re.match(pat, basename(f)).group(1) for f in ev_files]
    ev_names = map(lambda x: x.lower(), ev_names)
    ax.set_xticks(range(len(ev_names)))
    ax.set_yticks(range(len(ev_names)))
    ax.set_xticklabels(ev_names, fontsize=6, rotation="vertical")
    ax.set_yticklabels(ev_names, fontsize=6)
    out_file = abspath("design_correlation.png")
    plt.savefig(out_file)
    return out_file

def plot_residual(resid_file, background_file):
    """Plot the model residual. Use FSL for now; should move to pylab."""
    from os.path import abspath
    from nibabel import load
    from subprocess import call
    from scipy.stats import scoreatpercentile
    ov_nii = "sigmasquareds_overlay.nii.gz"
    out_file = abspath("sigmasquareds.png")
    err_img = load(resid_file)
    err_data = err_img.get_data()
    low = scoreatpercentile(err_data.ravel(), 2)
    high = scoreatpercentile(err_data.ravel(), 98)
    n_slice = err_img.shape[-1]
    native = n_slice < 50  # arbitrary but should work fine
    width = 750 if native else 872
    sample = 1 if native else 2
    del err_data
    call(["overlay", "1", "0", background_file, "-a",
          resid_file, str(low), str(high), ov_nii])
    call(["slicer", ov_nii, "-S", str(sample), str(width), out_file])
    return out_file

def plot_zstats(background_file, zstat_files, contrasts):
    """Plot the zstats and return a directory of pngs."""
    from os.path import abspath
    from nibabel import load
    from subprocess import call
    bg_img = load(background_file)
    n_slice = bg_img.shape[-1]
    native = n_slice < 50
    width = 750 if native else 872
    sample = 1 if native else 2
    out_files = []
    if not isinstance(zstat_files, list):
        zstat_files = [zstat_files]
    for i, contrast in enumerate(contrasts):
        zstat_file = zstat_files[i]
        zstat_png = "zstat%d.png" % (i + 1)
        ov_nii = "zstat%d_overlay.nii.gz" % (i + 1)
        call(["overlay", "1", "0", background_file, "-a",
              zstat_file, "2.3", "10",
              zstat_file, "-2.3", "-10", ov_nii])
        call(["slicer", ov_nii, "-S", str(sample), str(width), zstat_png])
        out_files.append(abspath(zstat_png))
    return out_files

def dump_exp_info(exp_info, timeseries):
    """Dump the exp_info dict into a json file."""
    from copy import deepcopy
    from os.path import abspath
    import json
    json_file = abspath("experiment_info.json")
    json_info = deepcopy(exp_info)
    del json_info["contrasts"]
    with open(json_file, "w") as fp:
        json.dump(json_info, fp, sort_keys=True, indent=4)
    return json_file

def write_model_report(subject_id, design_image, design_corr,
                       residual, zstat_pngs, contrast_names):
    """Write model report info to rst and convert to pdf/html."""
    import time
    from lyman.tools import write_workflow_report
    from lyman.workflows.reporting import model_report_template

    # Fill in the initial report template dict
    report_dict = dict(now=time.asctime(),
                       subject_id=subject_id,
                       design_image=design_image,
                       design_corr=design_corr,
                       residual=residual)

    # Add the zstat image templates and update the dict
    for i, con in enumerate(contrast_names, 1):
        report_dict["con%d_name" % i] = con
        report_dict["zstat%d_png" % i] = zstat_pngs[i - 1]
        header = "Zstat %d: %s" % (i, con)
        model_report_template = "\n".join(
            [model_report_template,
             header,
             "".join(["^" for l in header]),
             "",
             "".join([".. image:: %(zstat", str(i), "_png)s"]),
             "    :width: 6.5in",
             ""])

    out_files = write_workflow_report("model",
                                      model_report_template,
                                      report_dict)
    return out_files

# Smaller helper functions
# ========================

def run_indices(design_files):
    """Given a list of files, return a list of 1-based integers."""
    return range(len(design_files))
<|endoftext|>


<commit_before>#!/usr/bin/env python3

import socket
import sys
import shodan
import nmap
import argparse

s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.connect(("8.8.8.8", 80))
ipaddress = s.getsockname()[0]
s.close()
nma = nmap.PortScannerAsync()

# This is where we will send our data to the shodan api in order to see if there are any potential vulnerabilities.
def checkVulns(host, port, scandata):
    print('Vulnerability Scanner is a work in progress at the moment. ')
    SHODAN_API_KEY = "Your API Key"

    api = shodan.Shodan(SHODAN_API_KEY)
    #The idea here is to parse through the scandata and construct strings to pass to the exploit api.
    try:
        # Search Shodan
        results = api.exploits.search('PHP 5.6')

        # Show the results
        print('Results found: %s' % results['total'])
        for result in results['matches']:
            print('IP: %s' % result)
            #print(result['data'])
            print('')
    except shodan.APIError as e:
        print('Error: %s' % e)

def callback_result(host, scan_result):
    if int(scan_result['nmap']['scanstats']['uphosts']) >= 1:
        print("------")
        print(host, scan_result['scan'])
        # grab_banner(host, scan_result["scan"][host]["tcp"])
        # print("------")

def main():
    parser = argparse.ArgumentParser(
        description='''pysec-recon automates the recon process of security posture assessment. Requires python-nmap, shodan, and argparse. You can combine multiple options -sfv for instance''',
        epilog="""Author: Michael Diviesti @michael_atx on twitter""")
    parser.add_argument('-t', type=str, default=ipaddress,
                        help='Target IP Address. If you do not specify one, your current IP address [' + ipaddress + '] will be used')
    parser.add_argument('-s', action='store_true',
                        help='Subnet Scan: Scans all ip addresses in the subnet of the provided address')
    parser.add_argument('-f', action='store_true', help='Full scan: Scans internal and external ip')
    parser.add_argument('-v', action='store_true',
                        help='Vulnerability Scan: check for vulnerabilities in all found ips and ports')
    args = parser.parse_args()

    if args.s:
        targetlist = args.t.split(".")
        target = '.'.join(targetlist[:-1] + ["0/24"])
    else:
        target = args.t
    print('scanning ' + target)
    # if args.f:
    #     print(args.f)
    # if args.v:
    #     print(args.v)

    nma.scan(hosts=target, callback=callback_result)
    while nma.still_scanning():
        nma.wait(2)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)
<commit_msg>scanning both internal and external ip addresses work as well as the check against shodan's exploit api. Removed a bunch of print statements for now just to see how the vuln scanning works out but, will add some back tonight that will print out the NMAP responses. I might need to add a --verbose switch or something. There is quite a bit of code cleanup that needs to happen as well.<commit_after>#!/usr/bin/env python3

import socket
import sys
import shodan
import nmap
import argparse
import colorama
import time

s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.connect(("8.8.8.8", 80))
ipaddress = s.getsockname()[0]
s.close()
nma = nmap.PortScannerAsync()

SHODAN_API_KEY = "YOUR_SHODAN_API_KEY"
target = ''
subnet = False
full = False
vuln = False
external = ''

# This is where we will send our data to the shodan api in order to see if there are any potential vulnerabilities.
def checkVulns(host, scandata):
    global SHODAN_API_KEY
    for k, item in scandata.items():
        # print(item)
        if str(item['product']) != '' and str(item['version']) != '':
            strSearch = item['product'] + ' ' + item['version']

            # The idea here is to parse through the scandata and construct strings to pass to the exploit api.
            try:
                # Search Shodan
                api = shodan.Shodan(SHODAN_API_KEY)
                results = api.exploits.search(strSearch)

                # Show the results
                if results['total'] >= 1:
                    print(host + ' ' + str(k), colorama.Fore.RED + strSearch + colorama.Fore.RESET)
                else:
                    print(host + ' ' + str(k), colorama.Fore.GREEN + strSearch + colorama.Fore.RESET)
            except shodan.APIError as e:
                print(host + 'Search Term: ' + strSearch + ' Error: %s' % e)
        time.sleep(2)

def callback_result(host, scan_result):
    global vuln
    if int(scan_result['nmap']['scanstats']['uphosts']) >= 1:
        if vuln:
            if 'tcp' in scan_result['scan'][host]:
                checkVulns(host, scan_result['scan'][host]['tcp'])

                # print(host, scan_result['scan'])
                # grab_banner(host, scan_result["scan"][host]["tcp"])
                # print("------")

def doscan(host):
    nma.scan(hosts=host, callback=callback_result)
    while nma.still_scanning():
        nma.wait(2)

def main():
    global target
    global vuln
    global full
    global external
    global SHODAN_API_KEY
    parser = argparse.ArgumentParser(
        description='''pysec-recon automates the recon process of security posture assessment. Requires python-nmap, shodan, and argparse. You can combine multiple options -sfv for instance''',
        epilog="""Author: Michael Diviesti @michael_atx on twitter""")
    parser.add_argument('-t', type=str, default=ipaddress,
                        help='Target IP Address. If you do not specify one, your current IP address [' + ipaddress + '] will be used')
    parser.add_argument('-s', action='store_true',
                        help='Subnet Scan: Scans all ip addresses in the subnet of the provided address')
    parser.add_argument('-f', action='store_true', help='Full scan: Scans internal and external ip')
    parser.add_argument('-v', action='store_true',
                        help='Vulnerability Scan: check for vulnerabilities in all found ips and ports')
    args = parser.parse_args()

    if args.s:
        targetlist = args.t.split(".")
        target = '.'.join(targetlist[:-1] + ["0/24"])
    else:
        target = args.t
    print('scanning ' + target)

    if args.f:
        full = True
        api = shodan.Shodan(SHODAN_API_KEY)
        external = api.tools.myip()
    if args.v:
        vuln = True

    doscan(target)

    if full:
        doscan(external)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# pyswip.easy -- PySWIP helper functions
# (c) 2006-2007 Yüce TEKOL
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA 02110-1301, USA.

from pyswip.core import *

PYSWIP_UNIFY = 274700

class InvalidTypeError(TypeError):
    def __init__(self, *args):
        type = args and args[0] or "Unknown"
        msg = "Term is expected to be of type: '%s'" % type
        Exception.__init__(self, msg, *args)

class Atom(object):
    __slots__ = "handle","chars"

    def __init__(self, handleOrChars):
        """Create an atom.
        ``handleOrChars``: handle or string of the atom.
        """
        if isinstance(handleOrChars, basestring):
            self.handle = PL_new_atom(handleOrChars)
            self.chars = handleOrChars
        else:
            self.handle = handleOrChars
            PL_register_atom(self.handle)
            self.chars = c_char_p(PL_atom_chars(self.handle)).value

    def fromTerm(cls, term):
        """Create an atom from a Term or term handle."""
        if isinstance(term, Term):
            term = term.handle

        a = atom_t()
        if PL_get_atom(term, addressof(a)):
            return cls(a.value)

    fromTerm = classmethod(fromTerm)

    def __del__(self):
        PL_unregister_atom(self.handle)

    value = property(lambda s:s.chars)

    def __str__(self):
        if self.chars is not None:
            return self.chars
        else:
            return self.__repr__()

    def __repr__(self):
        return str(self.handle).join(["Atom('", "')"])

class Term(object):
    __slots__ = "handle","chars","__value"
    def __init__(self, handle=None):
        if handle:
            #self.handle = PL_copy_term_ref(handle)
            self.handle = handle
        else:
            self.handle = PL_new_term_ref()
        self.chars = None

    def get_value(self):
        pass

class Variable(object):
    __slots__ = "handle","chars"

    def __init__(self, handle=None, name=None):
        self.chars = None
        if name:
            self.chars = name
        if handle:
            self.handle = handle
            s = create_string_buffer("\00"*64)  # FIXME:
            ptr = cast(s, c_char_p)
            if PL_get_chars(handle, byref(ptr), CVT_VARIABLE|BUF_RING):
                self.chars = ptr.value
        #else:
        #    self.handle = PL_new_term_ref()
        #    PL_put_variable(self.handle)

    def unify(self, value):
        print "var unify"
        if type(value) == str:
            fun = PL_unify_atom_chars
        elif type(value) == int:
            fun = PL_unify_integer
        elif type(value) == bool:
            fun = PL_unify_bool
        elif type(value) == float:
            fun = PL_unify_float
        elif type(value) == list:
            fun = PL_unify_list
        else:
            raise

        t = PL_new_term_ref()
        fun(t, value)
        self.handle = t

    def get_value(self):
        return getTerm(self.handle)

    value = property(get_value, unify)

    def unified(self):
        return PL_term_type(self.handle) == PL_VARIABLE

    def __str__(self):
        if self.chars is not None:
            return self.chars
        else:
            return self.__repr__()

    def __repr__(self):
        return "Variable(%s)" % self.handle

    def put(self, term):
        PL_put_variable(term)
        self.handle = term

class Functor(object):
    __slots__ = "handle","name","arity","args","__value","a0"

    func = {}

    def __init__(self, handleOrName, arity=1, args=None, a0=None):
        """Create a functor.
        ``handleOrName``: functor handle, a string or an atom.
        """

        self.args = args or []
        self.arity = arity
        self.a0 = a0

        if isinstance(handleOrName, basestring):
            self.name = Atom(handleOrName)
            self.handle = PL_new_functor(self.name.handle, arity)
            self.__value = "Functor%d" % self.handle
        elif isinstance(handleOrName, Atom):
            self.name = handleOrName
            self.handle = PL_new_functor(self.name.handle, arity)
            self.__value = "Functor%d" % self.handle
        else:
            self.handle = handleOrName
            self.name = Atom(PL_functor_name(self.handle))
            self.arity = PL_functor_arity(self.handle)
            try:
                self.__value = self.func[self.handle](self.arity, *self.args)
            except KeyError:
                self.__value = "Functor%d" % self.handle

    def fromTerm(cls, term):
        """Create a functor from a Term or term handle."""
        if isinstance(term, Term):
            term = term.handle

        f = functor_t()
        if PL_get_functor(term, addressof(f)):
            # get args
            args = []
            arity = PL_functor_arity(f.value)
            # let's have all args be consecutive
            a0 = PL_new_term_refs(arity)
            for i, a in enumerate(range(1, arity + 1)):
                if PL_get_arg(a, term, a0 + i):
                    args.append(getTerm(a0 + i))

            return cls(f.value, args=args, a0=a0)

    fromTerm = classmethod(fromTerm)

    value = property(lambda s: s.__value)

    def __call__(self, *args):
        assert self.arity == len(args)
        a = PL_new_term_refs(len(args))
        #a = PL_new_term_ref()
        for i, arg in enumerate(args):
            putTerm(a + i, arg)
            #putTerm(a, arg)

        t = PL_new_term_ref()
        PL_cons_functor_v(t, self.handle, a)
        return Term(t)

    def __str__(self):
        if self.name is not None and self.arity is not None:
            return "%s(%d)" % (self.name,self.arity)
        else:
            return self.__repr__()

    def __repr__(self):
        return "".join(["Functor(", ",".join(str(x) for x in [self.handle,self.arity]+self.args), ")"])

def _unifier(arity, *args):
    assert arity == 2
    #if PL_is_variable(args[0]):
    #    args[0].unify(args[1])
    try:
        return {args[0].chars:args[1].value}
    except AttributeError:
        return {args[0].chars:args[1]}

Functor.func[274700] = _unifier

def putTerm(term, value):
    #print "putterm", term, value
    if isinstance(value, Term):
        PL_put_term(term, value.handle)
    elif isinstance(value, basestring):
        PL_put_atom_chars(term, value)
    elif isinstance(value, int):
        PL_put_integer(term, value)
    elif isinstance(value, Variable):
        value.put(term)
    elif isinstance(value, list):
        putList(term, value)
    elif isinstance(value, Atom):
        print "ATOM"
    elif isinstance(value, Functor):
        PL_put_functor(term, value.handle)
    else:
        raise Exception("Not implemented")

def putList(l, ls):
    h = PL_new_term_ref()
    PL_put_nil(h)
    a = PL_new_term_ref()  #PL_new_term_refs(len(ls))
    for item in reversed(ls):
        putTerm(a, item)
        PL_cons_list(h, a, h)
    PL_get_head(h, h)
    PL_unify(h, l)

# deprecated
def getAtomChars(t):
    """If t is an atom, return it as a string, otherwise raise InvalidTypeError.
    """
    s = c_char_p()
    if PL_get_atom_chars(t, addressof(s)):
        return s.value
    else:
        raise InvalidTypeError("atom")

def getAtom(t):
    """If t is an atom, return it , otherwise raise InvalidTypeError.
    """
    return Atom.fromTerm(t)

def getBool(t):
    """If t is of type bool, return it, otherwise raise InvalidTypeError.
    """
    b = c_int()
    if PL_get_long(t, addressof(b)):
        return bool(b.value)
    else:
        raise InvalidTypeError("bool")

def getLong(t):
    """If t is of type long, return it, otherwise raise InvalidTypeError.
    """
    i = c_long()
    if PL_get_long(t, addressof(i)):
        return i.value
    else:
        raise InvalidTypeError("long")

getInteger = getLong  # just an alias for getLong

def getFloat(t):
    """If t is of type float, return it, otherwise raise InvalidTypeError.
    """
    d = c_double()
    if PL_get_float(t, addressof(d)):
        return d.value
    else:
        raise InvalidTypeError("float")

def getString(t):
    """If t is of type string, return it, otherwise raise InvalidTypeError.
    """
    slen = c_int()
    s = c_char_p()
    if PL_get_string_chars(t, addressof(s), addressof(slen)):
        return s.value
    else:
        raise InvalidTypeError("string")

def getTerm(t):
    p = PL_term_type(t)
    if p < PL_TERM:
        return _getterm_router[p](t)
    elif PL_is_list(t):
        return getList(t)
    else:
        return getFunctor(t)

def getList(t):
    """Return t as a list.
    """
    head = PL_new_term_ref()
    result = []
    while PL_get_list(t, head, t):
        result.append(getTerm(head))

    return result

def getFunctor(t):
    """Return t as a functor
    """
    #f = functor_t()
    #if PL_get_functor(t, addressof(f)):
    #    return Functor(f.value, term=t)
    #else:
    #    raise InvalidTypeError("functor")
    return Functor.fromTerm(t)

def getVariable(t):
    return Variable(t)

_getterm_router = {
                    PL_VARIABLE:getVariable, PL_ATOM:getAtom, PL_STRING:getString,
                    PL_INTEGER:getInteger, PL_FLOAT:getFloat,
                    PL_TERM:getTerm
                    }

unifyInteger = PL_unify_integer

def _callbackWrapper(arity=1):
    return CFUNCTYPE(*([foreign_t] + [term_t]*arity))

def _foreignWrapper(fun):
    def wrapper(*args):
        args = [getTerm(arg) for arg in args]
        r = fun(*args)
        return (r is None) and True or r
    return wrapper

def registerForeign(func, name=None, arity=None, flags=0):
    """Register a Python predicate
    ``func``: Function to be registered. The function should return a value in
    ``foreign_t``, ``True`` or ``False``.
    ``name`` : Name of the function. If this value is not used, ``func.func_name``
    should exist.
    ``arity``: Arity (number of arguments) of the function. If this value is not
    used, ``func.arity`` should exist.
    """
    if arity is None:
        arity = func.arity

    if name is None:
        name = func.func_name

    return PL_register_foreign(name, arity,
            _callbackWrapper(arity)(_foreignWrapper(func)), flags)

newTermRef = PL_new_term_ref

def newTermRefs(count):
    a = PL_new_term_refs(count)
    return range(a, a + count)

def call(term, module=None):
    """Call term in module.
    ``term``: a Term or term handle
    """
    if isinstance(term, Term):
        term = term.handle

    return PL_call(term, module)

def newModule(name):
    """Create a new module.
    ``name``: An Atom or a string
    """
    if isinstance(name, basestring):
        name = Atom(name)

    return PL_new_module(name.handle)

class Query(object):
    qid = None
    fid = None
    comma = Functor(",", 2)

    def __init__(self, *terms, **kwargs):
        for key in kwargs:
            if key not in ["flags", "module"]:
                raise Exception("Invalid kwarg: %s" % key, key)

        flags = kwargs.get("flags", PL_Q_NODEBUG|PL_Q_CATCH_EXCEPTION)
        module = kwargs.get("module", None)

        comma = Query.comma
        t = terms[0]
        for tx in terms[1:]:
            t = comma(t, tx)

        f = Functor.fromTerm(t)
        p = PL_pred(f.handle, module)
        Query.qid = PL_open_query(module, flags, p, f.a0)

    def __del__(self):
        if Query.qid is not None:
            self.closeQuery()

    def nextSolution():
        return PL_next_solution(Query.qid)

    nextSolution = staticmethod(nextSolution)

    def cutQuery():
        PL_cut_query(Query.qid)

    cutQuery = staticmethod(cutQuery)

    def closeQuery():
        PL_close_query(Query.qid)

    closeQuery = staticmethod(closeQuery)

def _test():
    from pyswip.prolog import Prolog
    p = Prolog()

    assertz = Functor("assertz")
    a = Functor("a_")

    call(assertz(a(10)))
    call(assertz(a([1,2,3])))
    call(assertz(a(11)))

    X = Variable()

    q = Query(a(X))
    while q.nextSolution():
        print ">", X.value

if __name__ == "__main__":
    _test()

<commit_msg>removed unnecessary prints from easy<commit_after># -*- coding: utf-8 -*-

# pyswip.easy -- PySWIP helper functions
# (c) 2006-2007 Yüce TEKOL
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA 02110-1301, USA.

from pyswip.core import *

PYSWIP_UNIFY = 274700

class InvalidTypeError(TypeError):
    def __init__(self, *args):
        type = args and args[0] or "Unknown"
        msg = "Term is expected to be of type: '%s'" % type
        Exception.__init__(self, msg, *args)

class Atom(object):
    __slots__ = "handle","chars"

    def __init__(self, handleOrChars):
        """Create an atom.
        ``handleOrChars``: handle or string of the atom.
        """
        if isinstance(handleOrChars, basestring):
            self.handle = PL_new_atom(handleOrChars)
            self.chars = handleOrChars
        else:
            self.handle = handleOrChars
            PL_register_atom(self.handle)
            self.chars = c_char_p(PL_atom_chars(self.handle)).value

    def fromTerm(cls, term):
        """Create an atom from a Term or term handle."""
        if isinstance(term, Term):
            term = term.handle

        a = atom_t()
        if PL_get_atom(term, addressof(a)):
            return cls(a.value)

    fromTerm = classmethod(fromTerm)

    def __del__(self):
        PL_unregister_atom(self.handle)

    value = property(lambda s:s.chars)

    def __str__(self):
        if self.chars is not None:
            return self.chars
        else:
            return self.__repr__()

    def __repr__(self):
        return str(self.handle).join(["Atom('", "')"])

class Term(object):
    __slots__ = "handle","chars","__value"
    def __init__(self, handle=None):
        if handle:
            #self.handle = PL_copy_term_ref(handle)
            self.handle = handle
        else:
            self.handle = PL_new_term_ref()
        self.chars = None

    def get_value(self):
        pass

class Variable(object):
    __slots__ = "handle","chars"

    def __init__(self, handle=None, name=None):
        self.chars = None
        if name:
            self.chars = name
        if handle:
            self.handle = handle
            s = create_string_buffer("\00"*64)  # FIXME:
            ptr = cast(s, c_char_p)
            if PL_get_chars(handle, byref(ptr), CVT_VARIABLE|BUF_RING):
                self.chars = ptr.value
        #else:
        #    self.handle = PL_new_term_ref()
        #    PL_put_variable(self.handle)

    def unify(self, value):
        if type(value) == str:
            fun = PL_unify_atom_chars
        elif type(value) == int:
            fun = PL_unify_integer
        elif type(value) == bool:
            fun = PL_unify_bool
        elif type(value) == float:
            fun = PL_unify_float
        elif type(value) == list:
            fun = PL_unify_list
        else:
            raise

        t = PL_new_term_ref()
        fun(t, value)
        self.handle = t

    def get_value(self):
        return getTerm(self.handle)

    value = property(get_value, unify)

    def unified(self):
        return PL_term_type(self.handle) == PL_VARIABLE

    def __str__(self):
        if self.chars is not None:
            return self.chars
        else:
            return self.__repr__()

    def __repr__(self):
        return "Variable(%s)" % self.handle

    def put(self, term):
        PL_put_variable(term)
        self.handle = term

class Functor(object):
    __slots__ = "handle","name","arity","args","__value","a0"

    func = {}

    def __init__(self, handleOrName, arity=1, args=None, a0=None):
        """Create a functor.
        ``handleOrName``: functor handle, a string or an atom.
        """

        self.args = args or []
        self.arity = arity
        self.a0 = a0

        if isinstance(handleOrName, basestring):
            self.name = Atom(handleOrName)
            self.handle = PL_new_functor(self.name.handle, arity)
            self.__value = "Functor%d" % self.handle
        elif isinstance(handleOrName, Atom):
            self.name = handleOrName
            self.handle = PL_new_functor(self.name.handle, arity)
            self.__value = "Functor%d" % self.handle
        else:
            self.handle = handleOrName
            self.name = Atom(PL_functor_name(self.handle))
            self.arity = PL_functor_arity(self.handle)
            try:
                self.__value = self.func[self.handle](self.arity, *self.args)
            except KeyError:
                self.__value = "Functor%d" % self.handle

    def fromTerm(cls, term):
        """Create a functor from a Term or term handle."""
        if isinstance(term, Term):
            term = term.handle

        f = functor_t()
        if PL_get_functor(term, addressof(f)):
            # get args
            args = []
            arity = PL_functor_arity(f.value)
            # let's have all args be consecutive
            a0 = PL_new_term_refs(arity)
            for i, a in enumerate(range(1, arity + 1)):
                if PL_get_arg(a, term, a0 + i):
                    args.append(getTerm(a0 + i))

            return cls(f.value, args=args, a0=a0)

    fromTerm = classmethod(fromTerm)

    value = property(lambda s: s.__value)

    def __call__(self, *args):
        assert self.arity == len(args)
        a = PL_new_term_refs(len(args))
        #a = PL_new_term_ref()
        for i, arg in enumerate(args):
            putTerm(a + i, arg)
            #putTerm(a, arg)

        t = PL_new_term_ref()
        PL_cons_functor_v(t, self.handle, a)
        return Term(t)

    def __str__(self):
        if self.name is not None and self.arity is not None:
            return "%s(%d)" % (self.name,self.arity)
        else:
            return self.__repr__()

    def __repr__(self):
        return "".join(["Functor(", ",".join(str(x) for x in [self.handle,self.arity]+self.args), ")"])

def _unifier(arity, *args):
    assert arity == 2
    #if PL_is_variable(args[0]):
    #    args[0].unify(args[1])
    try:
        return {args[0].chars:args[1].value}
    except AttributeError:
        return {args[0].chars:args[1]}

Functor.func[274700] = _unifier

def putTerm(term, value):
    if isinstance(value, Term):
        PL_put_term(term, value.handle)
    elif isinstance(value, basestring):
        PL_put_atom_chars(term, value)
    elif isinstance(value, int):
        PL_put_integer(term, value)
    elif isinstance(value, Variable):
        value.put(term)
    elif isinstance(value, list):
        putList(term, value)
    elif isinstance(value, Atom):
        print "ATOM"
    elif isinstance(value, Functor):
        PL_put_functor(term, value.handle)
    else:
        raise Exception("Not implemented")

def putList(l, ls):
    h = PL_new_term_ref()
    PL_put_nil(h)
    a = PL_new_term_ref()  #PL_new_term_refs(len(ls))
    for item in reversed(ls):
        putTerm(a, item)
        PL_cons_list(h, a, h)
    PL_get_head(h, h)
    PL_unify(h, l)

# deprecated
def getAtomChars(t):
    """If t is an atom, return it as a string, otherwise raise InvalidTypeError.
    """
    s = c_char_p()
    if PL_get_atom_chars(t, addressof(s)):
        return s.value
    else:
        raise InvalidTypeError("atom")

def getAtom(t):
    """If t is an atom, return it , otherwise raise InvalidTypeError.
    """
    return Atom.fromTerm(t)

def getBool(t):
    """If t is of type bool, return it, otherwise raise InvalidTypeError.
    """
    b = c_int()
    if PL_get_long(t, addressof(b)):
        return bool(b.value)
    else:
        raise InvalidTypeError("bool")

def getLong(t):
    """If t is of type long, return it, otherwise raise InvalidTypeError.
    """
    i = c_long()
    if PL_get_long(t, addressof(i)):
        return i.value
    else:
        raise InvalidTypeError("long")

getInteger = getLong  # just an alias for getLong

def getFloat(t):
    """If t is of type float, return it, otherwise raise InvalidTypeError.
    """
    d = c_double()
    if PL_get_float(t, addressof(d)):
        return d.value
    else:
        raise InvalidTypeError("float")

def getString(t):
    """If t is of type string, return it, otherwise raise InvalidTypeError.
    """
    slen = c_int()
    s = c_char_p()
    if PL_get_string_chars(t, addressof(s), addressof(slen)):
        return s.value
    else:
        raise InvalidTypeError("string")

def getTerm(t):
    p = PL_term_type(t)
    if p < PL_TERM:
        return _getterm_router[p](t)
    elif PL_is_list(t):
        return getList(t)
    else:
        return getFunctor(t)

def getList(t):
    """Return t as a list.
    """
    head = PL_new_term_ref()
    result = []
    while PL_get_list(t, head, t):
        result.append(getTerm(head))

    return result

def getFunctor(t):
    """Return t as a functor
    """
    #f = functor_t()
    #if PL_get_functor(t, addressof(f)):
    #    return Functor(f.value, term=t)
    #else:
    #    raise InvalidTypeError("functor")
    return Functor.fromTerm(t)

def getVariable(t):
    return Variable(t)

_getterm_router = {
                    PL_VARIABLE:getVariable, PL_ATOM:getAtom, PL_STRING:getString,
                    PL_INTEGER:getInteger, PL_FLOAT:getFloat,
                    PL_TERM:getTerm
                    }

unifyInteger = PL_unify_integer

def _callbackWrapper(arity=1):
    return CFUNCTYPE(*([foreign_t] + [term_t]*arity))

def _foreignWrapper(fun):
    def wrapper(*args):
        args = [getTerm(arg) for arg in args]
        r = fun(*args)
        return (r is None) and True or r
    return wrapper

def registerForeign(func, name=None, arity=None, flags=0):
    """Register a Python predicate
    ``func``: Function to be registered. The function should return a value in
    ``foreign_t``, ``True`` or ``False``.
    ``name`` : Name of the function. If this value is not used, ``func.func_name``
    should exist.
    ``arity``: Arity (number of arguments) of the function. If this value is not
    used, ``func.arity`` should exist.
    """
    if arity is None:
        arity = func.arity

    if name is None:
        name = func.func_name

    return PL_register_foreign(name, arity,
            _callbackWrapper(arity)(_foreignWrapper(func)), flags)

newTermRef = PL_new_term_ref

def newTermRefs(count):
    a = PL_new_term_refs(count)
    return range(a, a + count)

def call(term, module=None):
    """Call term in module.
    ``term``: a Term or term handle
    """
    if isinstance(term, Term):
        term = term.handle

    return PL_call(term, module)

def newModule(name):
    """Create a new module.
    ``name``: An Atom or a string
    """
    if isinstance(name, basestring):
        name = Atom(name)

    return PL_new_module(name.handle)

class Query(object):
    qid = None
    fid = None
    comma = Functor(",", 2)

    def __init__(self, *terms, **kwargs):
        for key in kwargs:
            if key not in ["flags", "module"]:
                raise Exception("Invalid kwarg: %s" % key, key)

        flags = kwargs.get("flags", PL_Q_NODEBUG|PL_Q_CATCH_EXCEPTION)
        module = kwargs.get("module", None)

        comma = Query.comma
        t = terms[0]
        for tx in terms[1:]:
            t = comma(t, tx)

        f = Functor.fromTerm(t)
        p = PL_pred(f.handle, module)
        Query.qid = PL_open_query(module, flags, p, f.a0)

    def __del__(self):
        if Query.qid is not None:
            self.closeQuery()

    def nextSolution():
        return PL_next_solution(Query.qid)

    nextSolution = staticmethod(nextSolution)

    def cutQuery():
        PL_cut_query(Query.qid)

    cutQuery = staticmethod(cutQuery)

    def closeQuery():
        PL_close_query(Query.qid)

    closeQuery = staticmethod(closeQuery)

def _test():
    from pyswip.prolog import Prolog
    p = Prolog()

    assertz = Functor("assertz")
    a = Functor("a_")

    call(assertz(a(10)))
    call(assertz(a([1,2,3])))
    call(assertz(a(11)))

    X = Variable()

    q = Query(a(X))
    while q.nextSolution():
        print ">", X.value

if __name__ == "__main__":
    _test()

<|endoftext|>


<commit_before>import ctrl.bbb as ctrl

class Controller(ctrl.Controller):

    def __init__(self, *vargs, **kwargs):

        # Initialize controller
        super().__init__(*vargs, **kwargs)

    def __reset(self):

        # call super
        super().__reset()

        # add source: encoder1
        self.add_device('encoder1',
                        'ctrl.bbb.encoder', 'Encoder',
                        type = 'source',
                        outputs = ['encoder1'],
                        encoder = 2,
                        ratio = 60 * 35.557)

        # add source: encoder2
        self.add_device('encoder2',
                        'ctrl.bbb.encoder', 'Encoder',
                        type = 'source',
                        outputs = ['encoder2'],
                        encoder = 1,
                        ratio = - 60 * 35.557)

        # add source: imu
        self.add_device('imu',
                        'ctrl.bbb.imu', 'Inclinometer',
                        type = 'source',
                        enable = False,
                        outputs = ['theta','theta_dot'])

        # add sink: motor1
        self.add_device('motor1',
                        'ctrl.bbb.motor', 'Motor',
                        type = 'sink',
                        enable = True,
                        inputs = ['motor1'],
                        pwm_pin = 'P9_14',
                        dir_A = 'P9_12',
                        dir_B = 'P9_13',
                        enable_pin = 'P9_41')

        # add sink: motor2
        self.add_device('motor2',
                        'ctrl.bbb.motor', 'Motor',
                        type = 'sink',
                        enable = True,
                        inputs = ['motor2'],
                        pwm_pin = 'P8_13',
                        dir_B = 'P8_45',
                        dir_A = 'P8_46',
                        enable_pin = 'P9_41')

<commit_msg>Improved mpu9150 c interface<commit_after>import warnings
import time
import ctrl.bbb as ctrl

class Controller(ctrl.Controller):

    def __init__(self, *vargs, **kwargs):

        # Initialize controller
        super().__init__(*vargs, **kwargs)

    def __reset(self):

        # call super
        super().__reset()

        # add source: encoder1
        self.add_device('encoder1',
                        'ctrl.bbb.encoder', 'Encoder',
                        type = 'source',
                        outputs = ['encoder1'],
                        encoder = 2,
                        ratio = 60 * 35.557)

        # add source: encoder2
        self.add_device('encoder2',
                        'ctrl.bbb.encoder', 'Encoder',
                        type = 'source',
                        outputs = ['encoder2'],
                        encoder = 1,
                        ratio = - 60 * 35.557)

        # add source: imu
        self.add_device('imu',
                        'ctrl.bbb.imu', 'Inclinometer',
                        type = 'source',
                        enable = False,
                        outputs = ['theta','theta_dot'])

        # add sink: motor1
        self.add_device('motor1',
                        'ctrl.bbb.motor', 'Motor',
                        type = 'sink',
                        enable = True,
                        inputs = ['motor1'],
                        pwm_pin = 'P9_14',
                        dir_A = 'P9_12',
                        dir_B = 'P9_13',
                        enable_pin = 'P9_41')

        # add sink: motor2
        self.add_device('motor2',
                        'ctrl.bbb.motor', 'Motor',
                        type = 'sink',
                        enable = True,
                        inputs = ['motor2'],
                        pwm_pin = 'P8_13',
                        dir_B = 'P8_45',
                        dir_A = 'P8_46',
                        enable_pin = 'P9_41')

        # Initializing devices
        warnings.warn("> Initializing devices ...")
        self.start()
        time.sleep(2)
        self.stop()
        warnings.warn("> Done.")
<|endoftext|>