<commit_before>#!/usr/bin/env python
"""
scans IPv4 subnet for SSH servers on Port 22 or other server ports.
Useful for machines that don't/can't have NMAP installed (e.g. Windows),
  and device does not have Avahi server.
I wanted to make it as cross-platform as possible,
  where the user would have only basic Python installed (Windows)

Michael Hirsch, Ph.D.

Note: timeout value bare minimum is 0.15 seconds for LAN,
    suggest using higher values say 0.25 or 0.35.

"""
from time import time
import logging
from typing import Union, List
import socket
import ipaddress as ip
import concurrent.futures
import itertools

TIMEOUT = 0.3
PORT = 22

def run(port: int = PORT,
        service: str = '',
        timeout: float = TIMEOUT,
        baseip: Union[str, ip.IPv4Address] = None,
        debug: bool = False) -> List[ip.IPv4Address]:

    if not baseip:
        ownip = getLANip()
        print('own address', ownip)
    else:
        ownip = ip.ip_address(baseip)

    assert isinstance(ownip, ip.IPv4Address)

    net = netfromaddress(ownip)

    return scanhosts(net, port, service, timeout, debug)

# %% (1) get LAN IP of laptop
def getLANip() -> Union[ip.IPv4Address, ip.IPv6Address]:
    """ get IP of own interface
     ref: http://stackoverflow.com/a/23822431
    """
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    # don't use link local here (169.254.x.x) unless you have a specific need
        try:
            s.connect(('<broadcast>', 0))
        except OSError:
            s.connect(('8.8.8.8', 80))  # for BSD/Mac

        name = s.getsockname()[0]

    return ip.ip_address(name)

# %% (2) scan subnet for SSH servers
def isportopen(host: ip.IPv4Address, port: int, service: str,
               timeout: float = TIMEOUT,
               verbose: bool = True) -> bool:
    h = host.exploded

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.settimeout(timeout)  # seconds
        try:
            s.connect((h, port))
            b = s.recv(32)  # arbitrary number of bytes
# %% service decode (optional)
            ok = validateservice(service, h, b)
            if ok and verbose:
                print('found', service, 'on', host, 'port', port)

            return ok

        except (ConnectionRefusedError, socket.timeout, socket.error):
            logging.info('no connection to {} {}'.format(h, port))
            return False

def validateservice(service: str, h: str, b: bytes) -> bool:
    if not b:  # empty reply
        return False
# %% non-empty reply
    try:
        """
        splitlines is in case the ASCII/UTF8 response is less than 32 bytes,
        hoping server sends a \r\n
        """
        u = b.splitlines()[0].decode('utf-8')
        print('\n', u)
    except UnicodeDecodeError:
        """
        must not have been utf8 encoding..., maybe latin1 or something else..
        """
        print('\n', b)
        return False

# %% optional service validation
    val = True
    if service and service not in u.lower():
        val = False

    return val

def netfromaddress(addr: ip.IPv4Address) -> ip.IPv4Network:

    assert isinstance(addr, ip.IPv4Address)

    net = ip.ip_network(addr.exploded.rsplit('.', 1)[0]+'.0/24')

    assert isinstance(net, ip.IPv4Network)

    return net

# %% main loop
def scanhosts(net: ip.IPv4Network,
              port: int, service: str,
              timeout: float, debug: bool) -> List[ip.IPv4Address]:
    """
    loops over xxx.xxx.xxx.1-254
    IPv4 only.
    """

    assert isinstance(net, ip.IPv4Network)

    print('searching', net)

    hosts = list(net.hosts())

    if debug:
        servers = [h for h in hosts if isportopen(h, port, service, timeout)]
    else:
        with concurrent.futures.ThreadPoolExecutor(max_workers=100) as exc:
            portsopen = exc.map(isportopen, hosts,
                                itertools.repeat(port), itertools.repeat(service))
            servers = list(itertools.compress(hosts, portsopen))

    return servers

def main():
    from argparse import ArgumentParser
    p = ArgumentParser('scan for hosts with open port, without NMAP')
    p.add_argument('-p', '--port', help='single port to try',
                   default=PORT, type=int)
    p.add_argument('-s', '--service',
                   help='string to match to qualify detections', default='')
    p.add_argument('-t', '--timeout',
                   help='timeout to wait for server',
                   default=TIMEOUT, type=float)
    p.add_argument('-b', '--baseip',
                   help='set a specific subnet to scan')
    p.add_argument('--debug', help='run single-threaded for debugging',
                   action='store_true')
    P = p.parse_args()

    servers = run(P.port, P.service, P.timeout, P.baseip, P.debug)

    tic = time()

    print('\n*** RESULTS ***')
    print('found', len(servers), P.service, 'servers on port', P.port,
          'in {:.1f} seconds:'.format(time()-tic))
    print('\n'.join(map(str, servers)))

if __name__ == '__main__':
    main()
<commit_msg>utf8 ignore<commit_after>#!/usr/bin/env python
"""
scans IPv4 subnet for SSH servers on Port 22 or other server ports.
Useful for machines that don't/can't have NMAP installed (e.g. Windows),
  and device does not have Avahi server.
I wanted to make it as cross-platform as possible,
  where the user would have only basic Python installed (Windows)

Michael Hirsch, Ph.D.

Note: timeout value bare minimum is 0.15 seconds for LAN,
    suggest using higher values say 0.25 or 0.35.

"""
from time import time
import logging
from typing import Union, List
import socket
import ipaddress as ip
import concurrent.futures
import itertools

TIMEOUT = 0.3
PORT = 22

def run(port: int = PORT,
        service: str = '',
        timeout: float = TIMEOUT,
        baseip: Union[str, ip.IPv4Address] = None,
        debug: bool = False) -> List[ip.IPv4Address]:

    if not baseip:
        ownip = getLANip()
        print('own address', ownip)
    else:
        ownip = ip.ip_address(baseip)

    assert isinstance(ownip, ip.IPv4Address)

    net = netfromaddress(ownip)

    return scanhosts(net, port, service, timeout, debug)

# %% (1) get LAN IP of laptop
def getLANip() -> Union[ip.IPv4Address, ip.IPv6Address]:
    """ get IP of own interface
     ref: http://stackoverflow.com/a/23822431
    """
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    # don't use link local here (169.254.x.x) unless you have a specific need
        try:
            s.connect(('<broadcast>', 0))
        except OSError:
            s.connect(('8.8.8.8', 80))  # for BSD/Mac

        name = s.getsockname()[0]

    return ip.ip_address(name)

# %% (2) scan subnet for SSH servers
def isportopen(host: ip.IPv4Address, port: int, service: str,
               timeout: float = TIMEOUT,
               verbose: bool = True) -> bool:
    h = host.exploded

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.settimeout(timeout)  # seconds
        try:
            s.connect((h, port))
            b = s.recv(32)  # arbitrary number of bytes
# %% service decode (optional)
            ok = validateservice(service, h, b)
            if ok and verbose:
                print('found', service, 'on', host, 'port', port)

            return ok

        except (ConnectionRefusedError, socket.timeout, socket.error):
            logging.info('no connection to {} {}'.format(h, port))
            return False

def validateservice(service: str, h: str, b: bytes) -> bool:
    if not b:  # empty reply
        return False
# %% non-empty reply
    """
    splitlines is in case the ASCII/UTF8 response is less than 32 bytes,
    hoping server sends a \r\n
    """
    u = b.splitlines()[0].decode('utf-8', 'ignore')
    print('\n', u)

# %% optional service validation
    val = True
    if service and service not in u.lower():
        val = False

    return val

def netfromaddress(addr: ip.IPv4Address) -> ip.IPv4Network:

    assert isinstance(addr, ip.IPv4Address)

    net = ip.ip_network(addr.exploded.rsplit('.', 1)[0]+'.0/24')

    assert isinstance(net, ip.IPv4Network)

    return net

# %% main loop
def scanhosts(net: ip.IPv4Network,
              port: int, service: str,
              timeout: float, debug: bool) -> List[ip.IPv4Address]:
    """
    loops over xxx.xxx.xxx.1-254
    IPv4 only.
    """

    assert isinstance(net, ip.IPv4Network)

    print('searching', net)

    hosts = list(net.hosts())

    if debug:
        servers = [h for h in hosts if isportopen(h, port, service, timeout)]
    else:
        with concurrent.futures.ThreadPoolExecutor(max_workers=100) as exc:
            portsopen = exc.map(isportopen, hosts,
                                itertools.repeat(port), itertools.repeat(service))
            servers = list(itertools.compress(hosts, portsopen))

    return servers

def main():
    from argparse import ArgumentParser
    p = ArgumentParser('scan for hosts with open port, without NMAP')
    p.add_argument('-p', '--port', help='single port to try',
                   default=PORT, type=int)
    p.add_argument('-s', '--service',
                   help='string to match to qualify detections', default='')
    p.add_argument('-t', '--timeout',
                   help='timeout to wait for server',
                   default=TIMEOUT, type=float)
    p.add_argument('-b', '--baseip',
                   help='set a specific subnet to scan')
    p.add_argument('--debug', help='run single-threaded for debugging',
                   action='store_true')
    P = p.parse_args()

    servers = run(P.port, P.service, P.timeout, P.baseip, P.debug)

    tic = time()

    print('\n*** RESULTS ***')
    print('found', len(servers), P.service, 'servers on port', P.port,
          'in {:.1f} seconds:'.format(time()-tic))
    print('\n'.join(map(str, servers)))

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


<commit_before>from data_importers.management.commands import BaseXpressDemocracyClubCsvImporter

class Command(BaseXpressDemocracyClubCsvImporter):
    council_id = "E06000041"
    addresses_name = "parl.2019-12-12/Version 1/merged.tsv"
    stations_name = "parl.2019-12-12/Version 1/merged.tsv"
    elections = ["parl.2019-12-12"]
    csv_delimiter = "\t"
    allow_station_point_from_postcode = False

    def address_record_to_dict(self, record):
        rec = super().address_record_to_dict(record)
        uprn = record.property_urn.strip().lstrip("0")

        if uprn == "14056341":
            rec["postcode"] = "RG10 8QB"

        if uprn in ["14020309"]:
            rec["accept_suggestion"] = True

        return rec
<commit_msg>import script for wokingham (closes #3547)<commit_after>from data_importers.management.commands import BaseXpressDemocracyClubCsvImporter

class Command(BaseXpressDemocracyClubCsvImporter):
    council_id = "WOK"
    addresses_name = "2021-04-12T09:10:21.537410/Democracy_Club__06May2021.tsv"
    stations_name = "2021-04-12T09:10:21.537410/Democracy_Club__06May2021.tsv"
    elections = ["2021-05-06"]
    csv_delimiter = "\t"
    csv_encoding = "windows-1252"

    def address_record_to_dict(self, record):
        uprn = record.property_urn.strip().lstrip("0")

        if uprn in [
            "14008163",  # LOWER COTTAGE, WHITE HORSE LANE, FINCHAMPSTEAD, WOKINGHAM
            "10094223552",  # 4 SHEERLANDS ROAD, ARBORFIELD
            "14006244",  # FIELD VIEW, MOLE ROAD, SINDLESHAM, WOKINGHAM
            "14063909",  # MAYS HILL LODGE, BEECH HILL ROAD, SPENCERS WOOD, READING
            "10032929468",  # LANE END FARM HOUSE, SHINFIELD ROAD, SHINFIELD, READING
        ]:
            return None

        if record.addressline6 in [
            "RG2 9LG",
            "RG6 4AG",
            "RG7 1NL",
            "RG41 4EL",
            "RG10 9HN",
        ]:
            return None

        return super().address_record_to_dict(record)

    def station_record_to_dict(self, record):

        # Civic Offices Shute End Wokingham Berkshire RG40 1WH
        if record.polling_place_id == "2740":
            record = record._replace(polling_place_postcode="")

        return super().station_record_to_dict(record)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
"""
scans IPv4 subnet for SSH servers on Port 22.
Useful for machines that don't/can't have NMAP installed (e.g. Windows), and device does not have Avahi server.
Yes this can be done simply via Linux shell, but I wanted to make it as cross-platform as possible
where the user would have only basic Python installed (Windows)

Michael Hirsch

Note: if using Python < 3.3, you will need
pip install ipaddress

Note: timeout value bare minimum is 0.15 seconds for LAN, suggest using higher values say 0.25 or 0.35 if you can stand the wait 254*0.35 seconds

"""
from __future__ import division,unicode_literals
from six import PY2
from time import time
import socket
from ipaddress import ip_address,ip_network,IPv4Network
#
if PY2: ConnectionRefusedError = OSError
#%% (1) get LAN IP of laptop
def getLANip():
    """ get IP of own interface
     ref: http://stackoverflow.com/a/23822431
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    s.connect(('<broadcast>', 0))
    name= s.getsockname()[0]
    s.close()
    return ip_address(name)
#%% (2) scan subnet for SSH servers
def isportopen(host,port,service,timeout=0.25):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(timeout) #seconds
    try:
        s.connect((host.exploded,port))
        r=s.recv(32).decode('utf8') #arbitrary number of bytes
        print(r)
        s.close()
    except (ConnectionRefusedError,socket.timeout,socket.error):
        return

    if service in r.lower():
        return True

def netfromaddress(addr):
    addr = ip_address(addr) #in case it's string
    if addr.version == 4:
        return ip_network(ownip.exploded.rsplit('.',1)[0]+'.0/24')
    else: #addr.version ==6
        raise NotImplementedError('https://www.6net.org/publications/standards/draft-chown-v6ops-port-scanning-implications-00.txt')
#%% main loop
def scanhosts(net,port,service,timeout):
    """
    loops over xxx.xxx.xxx.1-254
    IPv4 only.
    """
    if not isinstance(net,IPv4Network):
        net = netfromaddress(net)

    if net.version != 4:
        raise NotImplementedError('https://www.6net.org/publications/standards/draft-chown-v6ops-port-scanning-implications-00.txt')

    servers = []
    print('searching {}'.format(net))
    for t,a in enumerate(net.hosts()):
        if isportopen(a,port,service):
            servers.append(a)
            print('found {} on {} port {}'.format(service,a,port))
        if not t % 20:
            print('{:.1f} % done, {} {} servers detected on port {}'.format(t/255*100.,len(servers),service,port))
    return servers

if __name__ == '__main__':
    from argparse import ArgumentParser
    p = ArgumentParser('scan for hosts with open port, without NMAP')
    p.add_argument('-p','--port',help='single port to try',default=22,type=int)
    p.add_argument('-s','--service',help='string to match to qualify detections',default='ssh')
    p.add_argument('-t','--timeout',help='timeout to wait for server',default=0.2,type=float)
    p.add_argument('-b','--baseip',help='instead of using own IP, set a specific subnet to scan')
    p = p.parse_args()

    tic = time()
    ownip = getLANip() if not p.baseip else p.baseip
    print('own address ' + str(ownip))
    servers = scanhosts(ownip,p.port,p.service,p.timeout)
    print('found {} {} server IPs in {:.1f} seconds: \n'.format(
                         len(servers),p.service,time()-tic)+str(servers))
<commit_msg>handle <32 byte replies, handle empty replies, 2&3 compatibility with utf8<commit_after>#!/usr/bin/env python3
"""
scans IPv4 subnet for SSH servers on Port 22.
Useful for machines that don't/can't have NMAP installed (e.g. Windows), and device does not have Avahi server.
Yes this can be done simply via Linux shell, but I wanted to make it as cross-platform as possible
where the user would have only basic Python installed (Windows)

Michael Hirsch

Note: if using Python < 3.3, you will need
pip install ipaddress

Note: timeout value bare minimum is 0.15 seconds for LAN, suggest using higher values say 0.25 or 0.35 if you can stand the wait 254*0.35 seconds

"""
from __future__ import division,unicode_literals
from six import PY2
from time import time
import socket
from ipaddress import ip_address,ip_network,IPv4Network
#
if PY2: ConnectionRefusedError = OSError
#%% (1) get LAN IP of laptop
def getLANip():
    """ get IP of own interface
     ref: http://stackoverflow.com/a/23822431
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    s.connect(('<broadcast>', 0))
    name= s.getsockname()[0]
    s.close()
    return ip_address(name.encode('utf-8').decode('utf-8')) #encode.decode is used for python2 and python3 compatibility
#%% (2) scan subnet for SSH servers
def isportopen(host,port,service,timeout=0.25):
    h = host.exploded
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(timeout) #seconds
    try:
        s.connect((h,port))
        b=s.recv(32) #arbitrary number of bytes

        if not b: #empty reply
            return

        try:
            u = b.splitlines()[0].decode('utf-8') #splitlines is in case the ASCII/UTF8 response is less than 32 bytes, hoping server sends a \r\n
        except UnicodeDecodeError: # must not have been utf8 encoding..., maybe latin1 or something else..
            u = b
        print(u)
        s.close()
    except (ConnectionRefusedError,socket.timeout,socket.error):
        return

    try:
        if service in u.lower():
            return True
    except UnicodeDecodeError:
        print('unable to decode server {} response'.format(h))
        return

def netfromaddress(addr):
    addr = ip_address(addr) #in case it's string
    if addr.version == 4:
        return ip_network(ownip.exploded.rsplit('.',1)[0]+'.0/24')
    else: #addr.version ==6
        raise NotImplementedError('https://www.6net.org/publications/standards/draft-chown-v6ops-port-scanning-implications-00.txt')
#%% main loop
def scanhosts(net,port,service,timeout):
    """
    loops over xxx.xxx.xxx.1-254
    IPv4 only.
    """
    if not isinstance(net,IPv4Network):
        net = netfromaddress(net)

    if net.version != 4:
        raise NotImplementedError('https://www.6net.org/publications/standards/draft-chown-v6ops-port-scanning-implications-00.txt')

    servers = []
    print('searching {}'.format(net))
    for t,a in enumerate(net.hosts()):
        if isportopen(a,port,service):
            servers.append(a)
            print('found {} on {} port {}'.format(service,a,port))
        if not t % 20:
            print('{:.1f} % done, {} {} servers detected on port {}'.format(t/255*100.,len(servers),service,port))
    return servers

if __name__ == '__main__':
    from argparse import ArgumentParser
    p = ArgumentParser('scan for hosts with open port, without NMAP')
    p.add_argument('-p','--port',help='single port to try',default=22,type=int)
    p.add_argument('-s','--service',help='string to match to qualify detections',default='ssh')
    p.add_argument('-t','--timeout',help='timeout to wait for server',default=0.2,type=float)
    p.add_argument('-b','--baseip',help='instead of using own IP, set a specific subnet to scan')
    p = p.parse_args()

    tic = time()
    ownip = getLANip() if not p.baseip else p.baseip
    print('own address ' + str(ownip))
    servers = scanhosts(ownip,p.port,p.service,p.timeout)
    print('\n*** RESULTS ***\n')
    print('found {} {} server IPs in {:.1f} seconds: \n'.format(
                         len(servers),p.service,time()-tic)+str(servers))
<|endoftext|>


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

from __future__ import unicode_literals

import re
from datetime import datetime, date

import six
import requests

__all__ = ('FioBank',)

str = six.text_type

def coerce_date(value):
    if isinstance(value, datetime):
        return value.date()
    elif isinstance(value, date):
        return value
    else:
        return datetime.strptime(value[:10], '%Y-%m-%d').date()

def sanitize_value(value, convert=None):
    if isinstance(value, six.string_types):
        value = value.strip() or None
    if convert and value is not None:
        return convert(value)
    return value

class FioBank(object):

    base_url = 'https://www.fio.cz/ib_api/rest/'

    actions = {
        'periods': 'periods/{token}/{from_date}/{to_date}/transactions.json',
        'by-id': 'by-id/{token}/{year}/{number}/transactions.json',
        'last': 'last/{token}/transactions.json',
        'set-last-id': 'set-last-id/{token}/{from_id}/',
        'set-last-date': 'set-last-date/{token}/{from_date}/',
    }

    # http://www.fio.cz/xsd/IBSchema.xsd
    transaction_schema = {
        'column22': ('transaction_id', str),
        'column0': ('date', coerce_date),
        'column1': ('amount', float),
        'column14': ('currency', str),
        'column2': ('account_number', str),
        'column10': ('account_name', str),
        'column3': ('bank_code', str),
        'column26': ('bic', str),
        'column12': ('bank_name', str),
        'column4': ('constant_symbol', str),
        'column5': ('variable_symbol', str),
        'column6': ('specific_symbol', str),
        'column7': ('user_identification', str),
        'column16': ('recipient_message', str),
        'column8': ('type', str),
        'column9': ('executor', str),
        'column18': ('specification', str),
        'column25': ('comment', str),
        'column17': ('instruction_id', str),
    }

    info_schema = {
        'accountid': ('account_number', str),
        'bankid': ('bank_code', str),
        'currency': ('currency', str),
        'iban': ('iban', str),
        'bic': ('bic', str),
        'closingbalance': ('balance', float),
    }

    _amount_re = re.compile(r'\-?\d+(\.\d+)? [A-Z]{3}')

    def __init__(self, token):
        self.token = token

    def _request(self, action, **params):
        template = self.base_url + self.actions[action]
        url = template.format(token=self.token, **params)

        response = requests.get(url)
        response.raise_for_status()

        if response.content:
            return response.json()
        return None

    def _parse_info(self, data):
        # parse data from API
        info = {}
        for key, value in data['accountStatement']['info'].items():
            key = key.lower()
            if key in self.info_schema:
                field_name, convert = self.info_schema[key]
                value = sanitize_value(value, convert)
                info[field_name] = value

        # make some refinements
        self._add_account_number_full(info)

        # return data
        return info

    def _parse_transactions(self, data):
        schema = self.transaction_schema
        try:
            entries = data['accountStatement']['transactionList']['transaction']  # NOQA
        except TypeError:
            entries = []

        for entry in entries:
            # parse entry from API
            trans = {}
            for column_name, column_data in entry.items():
                if not column_data:
                    continue
                field_name, convert = schema[column_name.lower()]
                value = sanitize_value(column_data['value'], convert)
                trans[field_name] = value

            # add missing fileds with None values
            for column_data_name, (field_name, convert) in schema.items():
                trans.setdefault(field_name, None)

            # make some refinements
            specification = trans.get('specification')
            is_amount = self._amount_re.match
            if specification is not None and is_amount(specification):
                amount, currency = trans['specification'].split(' ')
                trans['original_amount'] = float(amount)
                trans['original_currency'] = currency
            else:
                trans['original_amount'] = None
                trans['original_currency'] = None

            self._add_account_number_full(trans)

            # generate transaction data
            yield trans

    def _add_account_number_full(self, obj):
        account_number = obj.get('account_number')
        bank_code = obj.get('bank_code')

        if account_number is not None and bank_code is not None:
            account_number_full = '{}/{}'.format(account_number, bank_code)
        else:
            account_number_full = None

        obj['account_number_full'] = account_number_full

    def info(self):
        today = date.today()
        data = self._request('periods', from_date=today, to_date=today)
        return self._parse_info(data)

    def period(self, from_date, to_date):
        data = self._request('periods',
                             from_date=coerce_date(from_date),
                             to_date=coerce_date(to_date))
        return self._parse_transactions(data)

    def statement(self, year, number):
        data = self._request('by-id', year=year, number=number)
        return self._parse_transactions(data)

    def last(self, from_id=None, from_date=None):
        if from_id and from_date:
            raise ValueError('Only one constraint is allowed.')

        if from_id:
            self._request('set-last-id', from_id=from_id)
        elif from_date:
            self._request('set-last-date', from_date=coerce_date(from_date))

        return self._parse_transactions(self._request('last'))
<commit_msg>style: Polish according to pep8<commit_after># -*- coding: utf-8 -*-

from __future__ import unicode_literals

import re
from datetime import datetime, date

import six
import requests

__all__ = ('FioBank',)

str = six.text_type

def coerce_date(value):
    if isinstance(value, datetime):
        return value.date()
    elif isinstance(value, date):
        return value
    else:
        return datetime.strptime(value[:10], '%Y-%m-%d').date()

def sanitize_value(value, convert=None):
    if isinstance(value, six.string_types):
        value = value.strip() or None
    if convert and value is not None:
        return convert(value)
    return value

class ThrottlingError(Exception):
    """
        Throttling error raised when api is being used too fast.
    """
    def __str__(self):
        return 'Token should be used only once per 30s.'

class FioBank(object):

    base_url = 'https://www.fio.cz/ib_api/rest/'

    actions = {
        'periods': 'periods/{token}/{from_date}/{to_date}/transactions.json',
        'by-id': 'by-id/{token}/{year}/{number}/transactions.json',
        'last': 'last/{token}/transactions.json',
        'set-last-id': 'set-last-id/{token}/{from_id}/',
        'set-last-date': 'set-last-date/{token}/{from_date}/',
    }

    # http://www.fio.cz/xsd/IBSchema.xsd
    transaction_schema = {
        'column22': ('transaction_id', str),
        'column0': ('date', coerce_date),
        'column1': ('amount', float),
        'column14': ('currency', str),
        'column2': ('account_number', str),
        'column10': ('account_name', str),
        'column3': ('bank_code', str),
        'column26': ('bic', str),
        'column12': ('bank_name', str),
        'column4': ('constant_symbol', str),
        'column5': ('variable_symbol', str),
        'column6': ('specific_symbol', str),
        'column7': ('user_identification', str),
        'column16': ('recipient_message', str),
        'column8': ('type', str),
        'column9': ('executor', str),
        'column18': ('specification', str),
        'column25': ('comment', str),
        'column17': ('instruction_id', str),
    }

    info_schema = {
        'accountid': ('account_number', str),
        'bankid': ('bank_code', str),
        'currency': ('currency', str),
        'iban': ('iban', str),
        'bic': ('bic', str),
        'closingbalance': ('balance', float),
    }

    _amount_re = re.compile(r'\-?\d+(\.\d+)? [A-Z]{3}')

    def __init__(self, token):
        self.token = token

    def _request(self, action, **params):
        template = self.base_url + self.actions[action]
        url = template.format(token=self.token, **params)

        response = requests.get(url)
        response.raise_for_status()

        if response.content:
            return response.json()
        return None

    def _parse_info(self, data):
        # parse data from API
        info = {}
        for key, value in data['accountStatement']['info'].items():
            key = key.lower()
            if key in self.info_schema:
                field_name, convert = self.info_schema[key]
                value = sanitize_value(value, convert)
                info[field_name] = value

        # make some refinements
        self._add_account_number_full(info)

        # return data
        return info

    def _parse_transactions(self, data):
        schema = self.transaction_schema
        try:
            entries = data['accountStatement']['transactionList']['transaction']  # NOQA
        except TypeError:
            entries = []

        for entry in entries:
            # parse entry from API
            trans = {}
            for column_name, column_data in entry.items():
                if not column_data:
                    continue
                field_name, convert = schema[column_name.lower()]
                value = sanitize_value(column_data['value'], convert)
                trans[field_name] = value

            # add missing fileds with None values
            for column_data_name, (field_name, convert) in schema.items():
                trans.setdefault(field_name, None)

            # make some refinements
            specification = trans.get('specification')
            is_amount = self._amount_re.match
            if specification is not None and is_amount(specification):
                amount, currency = trans['specification'].split(' ')
                trans['original_amount'] = float(amount)
                trans['original_currency'] = currency
            else:
                trans['original_amount'] = None
                trans['original_currency'] = None

            self._add_account_number_full(trans)

            # generate transaction data
            yield trans

    def _add_account_number_full(self, obj):
        account_number = obj.get('account_number')
        bank_code = obj.get('bank_code')

        if account_number is not None and bank_code is not None:
            account_number_full = '{}/{}'.format(account_number, bank_code)
        else:
            account_number_full = None

        obj['account_number_full'] = account_number_full

    def info(self):
        today = date.today()
        data = self._request('periods', from_date=today, to_date=today)
        return self._parse_info(data)

    def period(self, from_date, to_date):
        data = self._request('periods',
                             from_date=coerce_date(from_date),
                             to_date=coerce_date(to_date))
        return self._parse_transactions(data)

    def statement(self, year, number):
        data = self._request('by-id', year=year, number=number)
        return self._parse_transactions(data)

    def last(self, from_id=None, from_date=None):
        if from_id and from_date:
            raise ValueError('Only one constraint is allowed.')

        if from_id:
            self._request('set-last-id', from_id=from_id)
        elif from_date:
            self._request('set-last-date', from_date=coerce_date(from_date))

        return self._parse_transactions(self._request('last'))
<|endoftext|>


<commit_before>import collections
import csv
import errno
import glob
import math
import multiprocessing
import os
import re
import shlex
import shutil
import subprocess
import sys
import StringIO

# 2014-12-23: For now, not using gevent, as it appears to conflict with python-daemon
## Import gevent and monkey-patch before importing bottle.
#from gevent import monkey
#monkey.patch_all()

from bottle import post, redirect, request, response, route, run, static_file, view, template

from wtforms import Form, BooleanField, IntegerField, RadioField, SelectField, StringField, validators, ValidationError

# Import matplotlib ourselves and make it use agg (not any GUI anything)
# before the analyze module pulls it in.
import matplotlib
matplotlib.use('Agg')
import analyze

_LOGDIR = "logs/"
_LOCKFILE = _LOGDIR + "explockfile"
_INIDIR = "ini/"
_IMGDIR = _LOGDIR + "img/"
_ARCHIVEDIR = _LOGDIR + "archive/"
_EXPSCRIPT = "python fishbox.py"

def _tracks():
    return sorted(glob.glob(_LOGDIR + "*-track.csv"))

def _inis():
    return sorted(glob.glob(_INIDIR + "*.ini"))

def _lock_exists():
    return os.path.isfile(_LOCKFILE)

def _imgs(name):
    return sorted(glob.glob(_IMGDIR + "%s*" % name))

def _mkdir(path):
    try:
        os.makedirs(path)
    except OSError as e:
        if e.errno == errno.EEXIST and os.path.isdir(path):
            # exists already, fine.
            pass
        else:
            raise

@route('/')
@view('index')
def index():
    tracks = []
    for index, track in enumerate(_tracks()):
        states = collections.Counter()
        heatmap = collections.Counter()
        with open(track) as f:
            i = -1  # so lines = 0 if file is empty
            for i, line in enumerate(f):
                vals = line.split(',')
                state, x, y = vals[1:4]
                if state == "init":
                    state = "lost"
                states[state] += 1
                try:
                    if state != "lost":
                        x = "%0.3f" % (float(x) - math.fmod(float(x), 1.0/15))
                        y = "%0.3f" % (float(y) - math.fmod(float(y), 1.0/10))
                        heatmap[(x,y)] += 1
                except ValueError:
                    pass  # not super important if we can't parse x,y
            lines = i+1
        name = track.split('/')[-1]
        aml = ["%0.3f" % (states[key] / float(lines)) for key in 'acquired', 'missing', 'lost']
        if heatmap:
            maxheat = max(heatmap.values())
            heat = [(key[0], key[1], str(float(value)/maxheat)) for key, value in heatmap.items()]
        else:
            heat = []
        tracks.append( (index, track, lines, aml, heat, _imgs(name)) )
    return dict(tracks=tracks)

@route('/new/')
def new_experiment():
    form = CreateExperimentForm()
    return template('new', form=form, lock_exists=_lock_exists())

@post('/delete_lockfile/')
def post_delete_lockfile():
    os.unlink(_LOCKFILE)

def _is_inifile(form, field):
    if field.data not in _inis():
        print str(field)
        raise ValidationError("Must specify an existing .ini file.")

def _name_is_sane(form, field):
    if re.search('\W', field.data):
        raise ValidationError("Experiment name must be alphanumeric characters only.")

class CreateExperimentForm(Form):
    ''' Form for creating a new experiment. '''
    expname = StringField("Experiment Name", [validators.Length(max=32), _name_is_sane])
    timelimit = IntegerField("Time Limit", [validators.NumberRange(min=1, max=24*60)])
    startfromtrig = BooleanField("startFromTrigger")
    stimulus = RadioField("Stimulus", choices=[('nostim', 'Off'), ('stim', 'On'), ('randstim', 'Randomized (equal chance off or on)')])
    inifile = SelectField(".ini File", choices=zip(_inis(), _inis()))

@post('/create/')
def post_create():
    if _lock_exists():
        return template('error', errormsg="It looks like an experiment is already running on this box.  Please wait for it to finish before starting another.")

    # validate form data
    form = CreateExperimentForm(request.forms)
    if not form.validate():
        return template('new', form=form, lock_exists=_lock_exists())

    expname = request.forms.expname
    timelimit = int(request.forms.timelimit)
    startfromtrig = bool(request.forms.startfromtrig)
    stimulus = request.forms.stimulus
    inifile = request.forms.inifile

    if sys.platform in ['cygwin', 'nt']:
        cmdparts = []
    else:
        cmdparts = ['sudo']  # fishbox.py must be run as root!
    cmdparts.append(_EXPSCRIPT)
    cmdparts.append("-t %d" % timelimit)
    if startfromtrig:
        cmdparts.append("--time-from-trigger")
    if stimulus == "nostim":
        cmdparts.append("--nostim")
    elif stimulus == "randstim":
        cmdparts.append("--randstim")
    cmdparts.append("--inifile %s" % inifile)
    cmdparts.append(expname)
    cmdline = ' '.join(cmdparts)
    cmdargs = shlex.split(cmdline)

    subprocess.Popen(cmdargs)

    redirect("/")

@route('/view/<logname:path>')
@view('view')
def view_log(logname):
    name = logname.split('/')[-1]
    return dict(imgs=_imgs(name), logname=logname)

@post('/stats/')
def post_stats():
    logs = request.query.logs.split('|')
    do_csv = request.query.csv
    stats = []
    all_keys = set()
    for log in logs:
        curstats = {}
        curstats['Log file'] = log
        g = analyze.Grapher()
        g.load(log)
        curstats.update(g.get_stats())
        for i in range(g.len_minutes):
            curstats.update(g.get_stats(minute=i))
            all_keys.update(curstats.keys())
        stats.append(curstats)
    for i in range(len(stats)):
        stat = stats[i]
        for k in all_keys:
            import numpy as np
            if k in stat:
                val = stat[k]
                if isinstance(val, (np.float32, np.float64)):
                    stat[k] = "%0.3f" % val
            else:
                stat[k] = ""

    all_keys.remove('Log file')  # will be added as first column
    all_keys = sorted(list(all_keys))
    all_keys[:0] = ['Log file']  # prepend 'Log file' header

    if do_csv:
        output = StringIO.StringIO()
        writer = csv.DictWriter(output, fieldnames=all_keys)
        writer.writeheader()
        for stat in stats:
            writer.writerow(stat)
        csvstring = output.getvalue()
        output.close()
        response.content_type = 'text/csv'
        response.headers['Content-Disposition'] = 'attachment; filename=fishystats.csv'
        return csvstring
    else:
        return template('stats', keys=all_keys, stats=stats)

def _do_analyze(logname):
    name = logname.split('/')[-1]
    g = analyze.Grapher()
    g.load(logname)
    g.plot()
    g.savefig(_IMGDIR + "%s.plot.png" % name)
    g.plot_heatmap()
    g.savefig(_IMGDIR + "%s.heat.png" % name)
    g.plot_heatmap(10)
    g.savefig(_IMGDIR + "%s.heat.10.png" % name)
    g.plot_leftright()
    g.savefig(_IMGDIR + "%s.leftright.png" % name)

@post('/analyze/')
def post_analyze():
    logname = request.query.path
    _do_analyze(logname)
    redirect("/view/%s" % logname)

@post('/analyze_all/')
def post_analyze_all():
    def do_all():
        for track in _tracks():
            _do_analyze(track)
    p = multiprocessing.Process(target=do_all)
    p.start()
    redirect("/")

@post('/compare/')
def post_compare():
    log1 = request.query.p1
    log2 = request.query.p2
    g1 = analyze.Grapher()
    g2 = analyze.Grapher()
    g1.load(log1)
    g2.load(log2)
    g1.plot_leftright()
    g2.plot_leftright(addplot=True)
    name1 = log1.split('/')[-1]
    name2 = log2.split('/')[-1]
    # XXX: bit of a hack doing pyplot stuff outside of Grapher...
    matplotlib.pyplot.legend([name1 + " Left", name2 + " Left", name1 + " Right", name2 + " Right"], fontsize=8, loc=2)

    imgname = _IMGDIR + "%s_%s_leftrights.png" % (name1, name2)
    g1.savefig(imgname)
    redirect("/" + imgname)

@post('/archive/')
def post_archive():
    logname = request.query.path
    name = logname.split('/')[-1].split('-track')[0]
    assert(name != '')
    allfiles = glob.glob(_LOGDIR + "%s*" % name)
    for f in allfiles:
        shutil.move(f, _ARCHIVEDIR)

@post('/unarchive/')
def post_unarchive():
    logname = request.query.path
    name = logname.split('/')[-1].split('-track')[0]
    allfiles = glob.glob(_ARCHIVEDIR + "%s*" % name)
    for f in allfiles:
        shutil.move(f, _LOGDIR)

@route('/logs/<filename:path>')
def static_logs(filename):
    if filename.endswith('.csv'):
        return static_file(filename, root=_LOGDIR, mimetype='text/plain')
    else:
        return static_file(filename, root=_LOGDIR)

@route('/ini/<filename:path>')
def static_inis(filename):
    return static_file(filename, root=_INIDIR)

if __name__ == '__main__':
    daemonize = False
    testing = False
    if len(sys.argv) > 1:
        arg1 = sys.argv[1]
        if arg1 == "--daemon":
            daemonize = True
        else:
            testing = True

    host = 'localhost' if testing else '0.0.0.0'

    # Create needed directories if not already there
    _mkdir(_IMGDIR)
    _mkdir(_ARCHIVEDIR)

    if daemonize:
        import daemon
        print("Launching daemon in the background.")
        logfile = os.path.join(
            os.getcwd(),
            "bottle.log"
        )
        with open(logfile, 'w+') as log:
            context = daemon.DaemonContext(
                working_directory=os.getcwd(),
                stdout=log,
                stderr=log
            )
            with context:
                # 2014-12-23: For now, not using gevent, as it appears to conflict with python-daemon
                #run(host=host, port=8080, server='gevent', debug=False, reloader=False)
                run(host=host, port=8080, debug=False, reloader=False)

    else:
        # 2014-12-23: For now, not using gevent, as it appears to conflict with python-daemon
        #run(host=host, port=8080, server='gevent', debug=testing, reloader=True)
        run(host=host, port=8080, debug=testing, reloader=True)
<commit_msg>Cache track data in memory for faster responses / less CPU waste.<commit_after>import collections
import csv
import errno
import glob
import math
import multiprocessing
import os
import re
import shlex
import shutil
import subprocess
import sys
import StringIO

# 2014-12-23: For now, not using gevent, as it appears to conflict with python-daemon
## Import gevent and monkey-patch before importing bottle.
#from gevent import monkey
#monkey.patch_all()

from bottle import post, redirect, request, response, route, run, static_file, view, template

from wtforms import Form, BooleanField, IntegerField, RadioField, SelectField, StringField, validators, ValidationError

# Import matplotlib ourselves and make it use agg (not any GUI anything)
# before the analyze module pulls it in.
import matplotlib
matplotlib.use('Agg')
import analyze

_LOGDIR = "logs/"
_LOCKFILE = _LOGDIR + "explockfile"
_INIDIR = "ini/"
_IMGDIR = _LOGDIR + "img/"
_ARCHIVEDIR = _LOGDIR + "archive/"
_EXPSCRIPT = "python fishbox.py"

def _tracks():
    return sorted(glob.glob(_LOGDIR + "*-track.csv"))

def _inis():
    return sorted(glob.glob(_INIDIR + "*.ini"))

def _lock_exists():
    return os.path.isfile(_LOCKFILE)

def _imgs(name):
    return sorted(glob.glob(_IMGDIR + "%s*" % name))

def _mkdir(path):
    try:
        os.makedirs(path)
    except OSError as e:
        if e.errno == errno.EEXIST and os.path.isdir(path):
            # exists already, fine.
            pass
        else:
            raise

def _get_track_data(track):
    states = collections.Counter()
    heatmap = collections.Counter()
    with open(track) as f:
        i = -1  # so lines = 0 if file is empty
        for i, line in enumerate(f):
            vals = line.split(',')
            state, x, y = vals[1:4]
            if state == "init":
                state = "lost"
            states[state] += 1
            try:
                if state != "lost":
                    x = "%0.3f" % (float(x) - math.fmod(float(x), 1.0/15))
                    y = "%0.3f" % (float(y) - math.fmod(float(y), 1.0/10))
                    heatmap[(x,y)] += 1
            except ValueError:
                pass  # not super important if we can't parse x,y
        lines = i+1
    aml = ["%0.3f" % (states[key] / float(lines)) for key in 'acquired', 'missing', 'lost']
    if heatmap:
        maxheat = max(heatmap.values())
        heat = [(key[0], key[1], str(float(value)/maxheat)) for key, value in heatmap.items()]
    else:
        heat = []

    return lines, aml, heat

# Cache computed a/m/l and heatmap data to avoid recomputation
track_data_cache = {}

@route('/')
@view('index')
def index():
    global track_data_cache
    tracks = []
    for index, track in enumerate(_tracks()):
        if track in track_data_cache:
            lines, aml, heat = track_data_cache[track]
        else:
            lines, aml, heat = _get_track_data(track)
            track_data_cache[track] = (lines, aml, heat)
        name = track.split('/')[-1]
        tracks.append( (index, track, lines, aml, heat, _imgs(name)) )
    return dict(tracks=tracks)

@route('/new/')
def new_experiment():
    form = CreateExperimentForm()
    return template('new', form=form, lock_exists=_lock_exists())

@post('/delete_lockfile/')
def post_delete_lockfile():
    os.unlink(_LOCKFILE)

def _is_inifile(form, field):
    if field.data not in _inis():
        print str(field)
        raise ValidationError("Must specify an existing .ini file.")

def _name_is_sane(form, field):
    if re.search('\W', field.data):
        raise ValidationError("Experiment name must be alphanumeric characters only.")

class CreateExperimentForm(Form):
    ''' Form for creating a new experiment. '''
    expname = StringField("Experiment Name", [validators.Length(max=32), _name_is_sane])
    timelimit = IntegerField("Time Limit", [validators.NumberRange(min=1, max=24*60)])
    startfromtrig = BooleanField("startFromTrigger")
    stimulus = RadioField("Stimulus", choices=[('nostim', 'Off'), ('stim', 'On'), ('randstim', 'Randomized (equal chance off or on)')])
    inifile = SelectField(".ini File", choices=zip(_inis(), _inis()))

@post('/create/')
def post_create():
    if _lock_exists():
        return template('error', errormsg="It looks like an experiment is already running on this box.  Please wait for it to finish before starting another.")

    # validate form data
    form = CreateExperimentForm(request.forms)
    if not form.validate():
        return template('new', form=form, lock_exists=_lock_exists())

    expname = request.forms.expname
    timelimit = int(request.forms.timelimit)
    startfromtrig = bool(request.forms.startfromtrig)
    stimulus = request.forms.stimulus
    inifile = request.forms.inifile

    if sys.platform in ['cygwin', 'nt']:
        cmdparts = []
    else:
        cmdparts = ['sudo']  # fishbox.py must be run as root!
    cmdparts.append(_EXPSCRIPT)
    cmdparts.append("-t %d" % timelimit)
    if startfromtrig:
        cmdparts.append("--time-from-trigger")
    if stimulus == "nostim":
        cmdparts.append("--nostim")
    elif stimulus == "randstim":
        cmdparts.append("--randstim")
    cmdparts.append("--inifile %s" % inifile)
    cmdparts.append(expname)
    cmdline = ' '.join(cmdparts)
    cmdargs = shlex.split(cmdline)

    subprocess.Popen(cmdargs)

    redirect("/")

@route('/view/<logname:path>')
@view('view')
def view_log(logname):
    name = logname.split('/')[-1]
    return dict(imgs=_imgs(name), logname=logname)

@post('/stats/')
def post_stats():
    logs = request.query.logs.split('|')
    do_csv = request.query.csv
    stats = []
    all_keys = set()
    for log in logs:
        curstats = {}
        curstats['Log file'] = log
        g = analyze.Grapher()
        g.load(log)
        curstats.update(g.get_stats())
        for i in range(g.len_minutes):
            curstats.update(g.get_stats(minute=i))
            all_keys.update(curstats.keys())
        stats.append(curstats)
    for i in range(len(stats)):
        stat = stats[i]
        for k in all_keys:
            import numpy as np
            if k in stat:
                val = stat[k]
                if isinstance(val, (np.float32, np.float64)):
                    stat[k] = "%0.3f" % val
            else:
                stat[k] = ""

    all_keys.remove('Log file')  # will be added as first column
    all_keys = sorted(list(all_keys))
    all_keys[:0] = ['Log file']  # prepend 'Log file' header

    if do_csv:
        output = StringIO.StringIO()
        writer = csv.DictWriter(output, fieldnames=all_keys)
        writer.writeheader()
        for stat in stats:
            writer.writerow(stat)
        csvstring = output.getvalue()
        output.close()
        response.content_type = 'text/csv'
        response.headers['Content-Disposition'] = 'attachment; filename=fishystats.csv'
        return csvstring
    else:
        return template('stats', keys=all_keys, stats=stats)

def _do_analyze(logname):
    name = logname.split('/')[-1]
    g = analyze.Grapher()
    g.load(logname)
    g.plot()
    g.savefig(_IMGDIR + "%s.plot.png" % name)
    g.plot_heatmap()
    g.savefig(_IMGDIR + "%s.heat.png" % name)
    g.plot_heatmap(10)
    g.savefig(_IMGDIR + "%s.heat.10.png" % name)
    g.plot_leftright()
    g.savefig(_IMGDIR + "%s.leftright.png" % name)

@post('/analyze/')
def post_analyze():
    logname = request.query.path
    _do_analyze(logname)
    redirect("/view/%s" % logname)

@post('/analyze_all/')
def post_analyze_all():
    def do_all():
        for track in _tracks():
            _do_analyze(track)
    p = multiprocessing.Process(target=do_all)
    p.start()
    redirect("/")

@post('/compare/')
def post_compare():
    log1 = request.query.p1
    log2 = request.query.p2
    g1 = analyze.Grapher()
    g2 = analyze.Grapher()
    g1.load(log1)
    g2.load(log2)
    g1.plot_leftright()
    g2.plot_leftright(addplot=True)
    name1 = log1.split('/')[-1]
    name2 = log2.split('/')[-1]
    # XXX: bit of a hack doing pyplot stuff outside of Grapher...
    matplotlib.pyplot.legend([name1 + " Left", name2 + " Left", name1 + " Right", name2 + " Right"], fontsize=8, loc=2)

    imgname = _IMGDIR + "%s_%s_leftrights.png" % (name1, name2)
    g1.savefig(imgname)
    redirect("/" + imgname)

@post('/archive/')
def post_archive():
    logname = request.query.path
    name = logname.split('/')[-1].split('-track')[0]
    assert(name != '')
    allfiles = glob.glob(_LOGDIR + "%s*" % name)
    for f in allfiles:
        shutil.move(f, _ARCHIVEDIR)

@post('/unarchive/')
def post_unarchive():
    logname = request.query.path
    name = logname.split('/')[-1].split('-track')[0]
    allfiles = glob.glob(_ARCHIVEDIR + "%s*" % name)
    for f in allfiles:
        shutil.move(f, _LOGDIR)

@route('/logs/<filename:path>')
def static_logs(filename):
    if filename.endswith('.csv'):
        return static_file(filename, root=_LOGDIR, mimetype='text/plain')
    else:
        return static_file(filename, root=_LOGDIR)

@route('/ini/<filename:path>')
def static_inis(filename):
    return static_file(filename, root=_INIDIR)

if __name__ == '__main__':
    daemonize = False
    testing = False
    if len(sys.argv) > 1:
        arg1 = sys.argv[1]
        if arg1 == "--daemon":
            daemonize = True
        else:
            testing = True

    host = 'localhost' if testing else '0.0.0.0'

    # Create needed directories if not already there
    _mkdir(_IMGDIR)
    _mkdir(_ARCHIVEDIR)

    if daemonize:
        import daemon
        print("Launching daemon in the background.")
        logfile = os.path.join(
            os.getcwd(),
            "bottle.log"
        )
        with open(logfile, 'w+') as log:
            context = daemon.DaemonContext(
                working_directory=os.getcwd(),
                stdout=log,
                stderr=log
            )
            with context:
                # 2014-12-23: For now, not using gevent, as it appears to conflict with python-daemon
                #run(host=host, port=8080, server='gevent', debug=False, reloader=False)
                run(host=host, port=8080, debug=False, reloader=False)

    else:
        # 2014-12-23: For now, not using gevent, as it appears to conflict with python-daemon
        #run(host=host, port=8080, server='gevent', debug=testing, reloader=True)
        run(host=host, port=8080, debug=testing, reloader=True)
<|endoftext|>


<commit_before>from __future__ import print_function, division
from copy import copy
from datetime import timedelta
import numpy as np
import pandas as pd
import nilmtk
from nilmtk.timeframegroup import TimeFrameGroup
from nilmtk.timeframe import TimeFrame
from neuralnilm.data.source import Sequence
from neuralnilm.utils import check_windows
from neuralnilm.data.activationssource import ActivationsSource
from neuralnilm.consts import DATA_FOLD_NAMES

import logging
logger = logging.getLogger(__name__)

class RealAggregateSource(ActivationsSource):
    """
    Attributes
    ----------
    mains : dict
        Structure example:
        {<train | unseen_appliances | unseen_activations_of_seen_appliances>: {
             <building_name>: pd.Series of raw data
        }}
    mains_good_sections : dict
        Same structure as `mains`.
    sections_with_no_target : dict
        Same structure as `mains`.
    all_gaps : dict of pd.DataFrames
        Each key is a fold name.
        Each DF has columns:
        building, gap, duration, p (short for 'probability')
        p is used by _get_sequence_without_target().
    """
    def __init__(self, activations, target_appliance,
                 seq_length, filename, windows, sample_period,
                 target_inclusion_prob=0.5,
                 uniform_prob_of_selecting_each_building=True,
                 allow_incomplete_target=True,
                 include_incomplete_target_in_output=True,
                 allow_multiple_target_activations_in_aggregate=False,
                 include_multiple_targets_in_output=False,
                 rng_seed=None):
        self.activations = copy(activations)
        self.target_appliance = target_appliance
        self.seq_length = seq_length
        self.filename = filename
        check_windows(windows)
        self.windows = windows
        self.sample_period = sample_period
        self.target_inclusion_prob = target_inclusion_prob
        self.uniform_prob_of_selecting_each_building = (
            uniform_prob_of_selecting_each_building)
        self.allow_incomplete_target = allow_incomplete_target
        self.include_incomplete_target_in_output = (
            include_incomplete_target_in_output)
        self.allow_multiple_target_activations_in_aggregate = (
            allow_multiple_target_activations_in_aggregate)
        self.include_multiple_targets_in_output = (
            include_multiple_targets_in_output)
        super(RealAggregateSource, self).__init__(rng_seed=rng_seed)

        self._load_mains_into_memory()
        self._remove_activations_with_no_mains()
        self._find_sections_with_no_target()
        self._compute_gap_probabilities()

    def _load_mains_into_memory(self):
        logger.info("Loading NILMTK mains...")

        # Load dataset
        dataset = nilmtk.DataSet(self.filename)

        self.mains = {}
        self.mains_good_sections = {}
        for fold, buildings_and_windows in self.windows.iteritems():
            for building_i, window in buildings_and_windows.iteritems():
                dataset.set_window(*window)
                elec = dataset.buildings[building_i].elec
                building_name = (
                    dataset.metadata['name'] +
                    '_building_{}'.format(building_i))

                logger.info(
                    "Loading mains for {}...".format(building_name))

                mains_meter = elec.mains()
                good_sections = mains_meter.good_sections()
                mains_data = mains_meter.power_series_all_data(
                    sample_period=self.sample_period,
                    sections=good_sections).dropna()

                def set_mains_data(dictionary, data):
                    dictionary.setdefault(fold, {})[building_name] = data

                if not mains_data.empty:
                    set_mains_data(self.mains, mains_data)
                    set_mains_data(self.mains_good_sections, good_sections)

                logger.info(
                    "Loaded mains data from building {} for fold {}"
                    " from {} to {}."
                    .format(building_name, fold,
                            mains_data.index[0], mains_data.index[-1]))

        dataset.store.close()
        logger.info("Done loading NILMTK mains data.")

    def _find_sections_with_no_target(self):
        """Finds the intersections of the mains good sections with the gaps
        between target appliance activations.
        """
        self.sections_with_no_target = {}
        seq_length_secs = self.seq_length * self.sample_period
        for fold, sects_per_building in self.mains_good_sections.iteritems():
            for building, good_sections in sects_per_building.iteritems():
                activations = (
                    self.activations[fold][self.target_appliance][building])
                mains = self.mains[fold][building]
                mains_good_sections = self.mains_good_sections[fold][building]
                gaps_between_activations = TimeFrameGroup()
                prev_end = mains.index[0]
                for activation in activations:
                    gap = TimeFrame(prev_end, activation.index[0])
                    gaps_between_activations.append(gap)
                    prev_end = activation.index[-1]
                gap = TimeFrame(prev_end, mains.index[-1])
                gaps_between_activations.append(gap)
                intersection = (
                    gaps_between_activations.intersection(mains_good_sections))
                intersection = intersection.remove_shorter_than(
                    seq_length_secs)
                self.sections_with_no_target.setdefault(fold, {})[building] = (
                    intersection)
                logger.info("Found {} sections without target for {} {}."
                            .format(len(intersection), fold, building))

    def _compute_gap_probabilities(self):
        # Choose a building and a gap
        self.all_gaps = {}
        for fold in DATA_FOLD_NAMES:
            all_gaps_for_fold = []
            for building, gaps in self.sections_with_no_target[fold].iteritems():
                gaps_for_building = [
                    (building, gap, gap.timedelta.total_seconds())
                    for gap in gaps]
                all_gaps_for_fold.extend(gaps_for_building)
            gaps_df = pd.DataFrame(
                all_gaps_for_fold, columns=['building', 'gap', 'duration'])
            gaps_df['p'] = gaps_df['duration'] / gaps_df['duration'].sum()
            self.all_gaps[fold] = gaps_df

    def _get_sequence_without_target(self, fold):
        # Choose a building and a gap
        all_gaps_for_fold = self.all_gaps[fold]
        n = len(all_gaps_for_fold)
        gap_i = self.rng.choice(n, p=all_gaps_for_fold['p'])
        row = all_gaps_for_fold.iloc[gap_i]
        building, gap = row['building'], row['gap']

        # Choose a start point in the gap
        latest_start_time = gap.end - timedelta(
            seconds=self.seq_length * self.sample_period)
        max_offset_seconds = (latest_start_time - gap.start).total_seconds()
        if max_offset_seconds <= 0:
            offset = 0
        else:
            offset = self.rng.randint(max_offset_seconds)
        start_time = gap.start + timedelta(seconds=offset)
        end_time = start_time + timedelta(
            seconds=(self.seq_length + 1) * self.sample_period)
        mains = self.mains[fold][building][start_time:end_time]
        seq = Sequence(self.seq_length)
        seq.input = mains.values[:self.seq_length]
        assert len(seq.input) == self.seq_length
        return seq

    def _has_sufficient_samples(self, data, start, end, threshold=0.8):
        if len(data) < 2:
            return False
        num_expected_samples = (
            (end - start).total_seconds() / self.sample_period)
        hit_rate = len(data) / num_expected_samples
        return (hit_rate >= threshold)

    def _remove_activations_with_no_mains(self):
        # First remove any activations where there is no mains data at all
        for fold, activations_for_appliance in self.activations.iteritems():
            activations_for_buildings = activations_for_appliance[
                self.target_appliance]
            buildings_to_remove = []
            for building in activations_for_buildings:
                mains_for_fold = self.mains[fold]
                if (building not in mains_for_fold and
                        building not in buildings_to_remove):
                    buildings_to_remove.append(building)
            for building in buildings_to_remove:
                self.activations[fold][self.target_appliance].pop(building)

        # Now check for places where mains has insufficient samples,
        # for example because the mains series has a break in it.
        for fold, activations_for_appliance in self.activations.iteritems():
            activations_for_buildings = activations_for_appliance[
                self.target_appliance]
            for building, activations in activations_for_buildings.iteritems():
                mains = self.mains[fold][building]
                activations_to_remove = []
                for i, activation in enumerate(activations):
                    activation_duration = (
                        activation.index[-1] - activation.index[0])
                    start = activation.index[0] - activation_duration
                    end = activation.index[-1] + activation_duration
                    mains_for_activ = mains[start:end]
                    if (start < mains.index[0] or
                            end > mains.index[-1] or not
                            self._has_sufficient_samples(
                                mains_for_activ, start, end)):
                        activations_to_remove.append(i)
                if activations_to_remove:
                    logger.info(
                        "Removing {} activations from fold '{}' building '{}'"
                        " because there was not enough mains data for"
                        " these activations. This leaves {} activations."
                        .format(
                            len(activations_to_remove), fold, building,
                            len(activations) - len(activations_to_remove)))
                activations_to_remove.reverse()
                for i in activations_to_remove:
                    activations.pop(i)
                self.activations[fold][self.target_appliance][building] = (
                    activations)

    def _get_sequence(self, fold='train', enable_all_appliances=False):
        if enable_all_appliances:
            raise ValueError("`enable_all_appliances` is not implemented yet"
                             " for RealAggregateSource!")

        if self.rng.binomial(n=1, p=self.target_inclusion_prob):
            seq = self._get_sequence_which_includes_target(fold=fold)
        else:
            seq = self._get_sequence_without_target(fold=fold)

        seq.input = seq.input[:, np.newaxis]
        seq.target = seq.target[:, np.newaxis]
        assert len(seq.input) == self.seq_length
        assert len(seq.target) == self.seq_length
        return seq

    def _get_sequence_which_includes_target(self, fold):
        MAX_RETRIES = 50
        for retry_i in range(MAX_RETRIES):
            seq = Sequence(self.seq_length)
            building_name = self._select_building(fold, self.target_appliance)
            activations = (
                self.activations[fold][self.target_appliance][building_name])
            activation_i = self._select_activation(activations)
            activation = activations[activation_i]
            positioned_activation, is_complete = (
                self._position_activation(
                    activation, is_target_appliance=True))
            if is_complete or self.include_incomplete_target_in_output:
                seq.target = positioned_activation
            else:
                seq.target = pd.Series(0, index=positioned_activation.index)

            # Check neighbouring activations
            mains_start = positioned_activation.index[0]
            mains_end = positioned_activation.index[-1]

            def neighbours_ok(neighbour_indicies):
                for i in neighbour_indicies:
                    activation = activations[i]
                    activation_duration = (
                        activation.index[-1] - activation.index[0])
                    neighbouring_activation_is_inside_mains_window = (
                        activation.index[0] >
                        (mains_start - activation_duration)
                        and activation.index[0] < mains_end)

                    if neighbouring_activation_is_inside_mains_window:
                        if self.allow_multiple_target_activations_in_aggregate:
                            if self.include_multiple_targets_in_output:
                                sum_target = seq.target.add(
                                    activation, fill_value=0)
                                is_complete = (
                                    sum_target.index == seq.target.index)
                                if self.allow_incomplete_target or is_complete:
                                    seq.target = sum_target[seq.target.index]
                        else:
                            return False  # need to retry
                    else:
                        return True  # everything checks out OK so far
                return True

            # Check forwards
            if not neighbours_ok(range(activation_i+1, len(activations))):
                continue
            # Check backwards
            if not neighbours_ok(range(activation_i-1, -1, -1)):
                continue

            # Get mains
            mains_for_building = self.mains[fold][activation.building]
            # load some additional data to make sure we have enough samples
            mains_end_extended = mains_end + timedelta(
                seconds=self.sample_period * 2)
            mains = mains_for_building[mains_start:mains_end_extended]
            seq.input = mains.values[:self.seq_length]
            if len(seq.input) != self.seq_length:
                continue
            return seq
        raise RuntimeError("No valid sequences found after {} retries!"
                           .format(MAX_RETRIES))

    @classmethod
    def _attrs_to_remove_for_report(cls):
        return [
            'activations', 'rng', 'mains', 'mains_good_sections',
            'sections_with_no_target']
<commit_msg>don't send 'all_gaps' to mongo!<commit_after>from __future__ import print_function, division
from copy import copy
from datetime import timedelta
import numpy as np
import pandas as pd
import nilmtk
from nilmtk.timeframegroup import TimeFrameGroup
from nilmtk.timeframe import TimeFrame
from neuralnilm.data.source import Sequence
from neuralnilm.utils import check_windows
from neuralnilm.data.activationssource import ActivationsSource
from neuralnilm.consts import DATA_FOLD_NAMES

import logging
logger = logging.getLogger(__name__)

class RealAggregateSource(ActivationsSource):
    """
    Attributes
    ----------
    mains : dict
        Structure example:
        {<train | unseen_appliances | unseen_activations_of_seen_appliances>: {
             <building_name>: pd.Series of raw data
        }}
    mains_good_sections : dict
        Same structure as `mains`.
    sections_with_no_target : dict
        Same structure as `mains`.
    all_gaps : dict of pd.DataFrames
        Each key is a fold name.
        Each DF has columns:
        building, gap, duration, p (short for 'probability')
        p is used by _get_sequence_without_target().
    """
    def __init__(self, activations, target_appliance,
                 seq_length, filename, windows, sample_period,
                 target_inclusion_prob=0.5,
                 uniform_prob_of_selecting_each_building=True,
                 allow_incomplete_target=True,
                 include_incomplete_target_in_output=True,
                 allow_multiple_target_activations_in_aggregate=False,
                 include_multiple_targets_in_output=False,
                 rng_seed=None):
        self.activations = copy(activations)
        self.target_appliance = target_appliance
        self.seq_length = seq_length
        self.filename = filename
        check_windows(windows)
        self.windows = windows
        self.sample_period = sample_period
        self.target_inclusion_prob = target_inclusion_prob
        self.uniform_prob_of_selecting_each_building = (
            uniform_prob_of_selecting_each_building)
        self.allow_incomplete_target = allow_incomplete_target
        self.include_incomplete_target_in_output = (
            include_incomplete_target_in_output)
        self.allow_multiple_target_activations_in_aggregate = (
            allow_multiple_target_activations_in_aggregate)
        self.include_multiple_targets_in_output = (
            include_multiple_targets_in_output)
        super(RealAggregateSource, self).__init__(rng_seed=rng_seed)

        self._load_mains_into_memory()
        self._remove_activations_with_no_mains()
        self._find_sections_with_no_target()
        self._compute_gap_probabilities()

    def _load_mains_into_memory(self):
        logger.info("Loading NILMTK mains...")

        # Load dataset
        dataset = nilmtk.DataSet(self.filename)

        self.mains = {}
        self.mains_good_sections = {}
        for fold, buildings_and_windows in self.windows.iteritems():
            for building_i, window in buildings_and_windows.iteritems():
                dataset.set_window(*window)
                elec = dataset.buildings[building_i].elec
                building_name = (
                    dataset.metadata['name'] +
                    '_building_{}'.format(building_i))

                logger.info(
                    "Loading mains for {}...".format(building_name))

                mains_meter = elec.mains()
                good_sections = mains_meter.good_sections()
                mains_data = mains_meter.power_series_all_data(
                    sample_period=self.sample_period,
                    sections=good_sections).dropna()

                def set_mains_data(dictionary, data):
                    dictionary.setdefault(fold, {})[building_name] = data

                if not mains_data.empty:
                    set_mains_data(self.mains, mains_data)
                    set_mains_data(self.mains_good_sections, good_sections)

                logger.info(
                    "Loaded mains data from building {} for fold {}"
                    " from {} to {}."
                    .format(building_name, fold,
                            mains_data.index[0], mains_data.index[-1]))

        dataset.store.close()
        logger.info("Done loading NILMTK mains data.")

    def _find_sections_with_no_target(self):
        """Finds the intersections of the mains good sections with the gaps
        between target appliance activations.
        """
        self.sections_with_no_target = {}
        seq_length_secs = self.seq_length * self.sample_period
        for fold, sects_per_building in self.mains_good_sections.iteritems():
            for building, good_sections in sects_per_building.iteritems():
                activations = (
                    self.activations[fold][self.target_appliance][building])
                mains = self.mains[fold][building]
                mains_good_sections = self.mains_good_sections[fold][building]
                gaps_between_activations = TimeFrameGroup()
                prev_end = mains.index[0]
                for activation in activations:
                    gap = TimeFrame(prev_end, activation.index[0])
                    gaps_between_activations.append(gap)
                    prev_end = activation.index[-1]
                gap = TimeFrame(prev_end, mains.index[-1])
                gaps_between_activations.append(gap)
                intersection = (
                    gaps_between_activations.intersection(mains_good_sections))
                intersection = intersection.remove_shorter_than(
                    seq_length_secs)
                self.sections_with_no_target.setdefault(fold, {})[building] = (
                    intersection)
                logger.info("Found {} sections without target for {} {}."
                            .format(len(intersection), fold, building))

    def _compute_gap_probabilities(self):
        # Choose a building and a gap
        self.all_gaps = {}
        for fold in DATA_FOLD_NAMES:
            all_gaps_for_fold = []
            for building, gaps in self.sections_with_no_target[fold].iteritems():
                gaps_for_building = [
                    (building, gap, gap.timedelta.total_seconds())
                    for gap in gaps]
                all_gaps_for_fold.extend(gaps_for_building)
            gaps_df = pd.DataFrame(
                all_gaps_for_fold, columns=['building', 'gap', 'duration'])
            gaps_df['p'] = gaps_df['duration'] / gaps_df['duration'].sum()
            self.all_gaps[fold] = gaps_df

    def _get_sequence_without_target(self, fold):
        # Choose a building and a gap
        all_gaps_for_fold = self.all_gaps[fold]
        n = len(all_gaps_for_fold)
        gap_i = self.rng.choice(n, p=all_gaps_for_fold['p'])
        row = all_gaps_for_fold.iloc[gap_i]
        building, gap = row['building'], row['gap']

        # Choose a start point in the gap
        latest_start_time = gap.end - timedelta(
            seconds=self.seq_length * self.sample_period)
        max_offset_seconds = (latest_start_time - gap.start).total_seconds()
        if max_offset_seconds <= 0:
            offset = 0
        else:
            offset = self.rng.randint(max_offset_seconds)
        start_time = gap.start + timedelta(seconds=offset)
        end_time = start_time + timedelta(
            seconds=(self.seq_length + 1) * self.sample_period)
        mains = self.mains[fold][building][start_time:end_time]
        seq = Sequence(self.seq_length)
        seq.input = mains.values[:self.seq_length]
        assert len(seq.input) == self.seq_length
        return seq

    def _has_sufficient_samples(self, data, start, end, threshold=0.8):
        if len(data) < 2:
            return False
        num_expected_samples = (
            (end - start).total_seconds() / self.sample_period)
        hit_rate = len(data) / num_expected_samples
        return (hit_rate >= threshold)

    def _remove_activations_with_no_mains(self):
        # First remove any activations where there is no mains data at all
        for fold, activations_for_appliance in self.activations.iteritems():
            activations_for_buildings = activations_for_appliance[
                self.target_appliance]
            buildings_to_remove = []
            for building in activations_for_buildings:
                mains_for_fold = self.mains[fold]
                if (building not in mains_for_fold and
                        building not in buildings_to_remove):
                    buildings_to_remove.append(building)
            for building in buildings_to_remove:
                self.activations[fold][self.target_appliance].pop(building)

        # Now check for places where mains has insufficient samples,
        # for example because the mains series has a break in it.
        for fold, activations_for_appliance in self.activations.iteritems():
            activations_for_buildings = activations_for_appliance[
                self.target_appliance]
            for building, activations in activations_for_buildings.iteritems():
                mains = self.mains[fold][building]
                activations_to_remove = []
                for i, activation in enumerate(activations):
                    activation_duration = (
                        activation.index[-1] - activation.index[0])
                    start = activation.index[0] - activation_duration
                    end = activation.index[-1] + activation_duration
                    mains_for_activ = mains[start:end]
                    if (start < mains.index[0] or
                            end > mains.index[-1] or not
                            self._has_sufficient_samples(
                                mains_for_activ, start, end)):
                        activations_to_remove.append(i)
                if activations_to_remove:
                    logger.info(
                        "Removing {} activations from fold '{}' building '{}'"
                        " because there was not enough mains data for"
                        " these activations. This leaves {} activations."
                        .format(
                            len(activations_to_remove), fold, building,
                            len(activations) - len(activations_to_remove)))
                activations_to_remove.reverse()
                for i in activations_to_remove:
                    activations.pop(i)
                self.activations[fold][self.target_appliance][building] = (
                    activations)

    def _get_sequence(self, fold='train', enable_all_appliances=False):
        if enable_all_appliances:
            raise ValueError("`enable_all_appliances` is not implemented yet"
                             " for RealAggregateSource!")

        if self.rng.binomial(n=1, p=self.target_inclusion_prob):
            seq = self._get_sequence_which_includes_target(fold=fold)
        else:
            seq = self._get_sequence_without_target(fold=fold)

        seq.input = seq.input[:, np.newaxis]
        seq.target = seq.target[:, np.newaxis]
        assert len(seq.input) == self.seq_length
        assert len(seq.target) == self.seq_length
        return seq

    def _get_sequence_which_includes_target(self, fold):
        MAX_RETRIES = 50
        for retry_i in range(MAX_RETRIES):
            seq = Sequence(self.seq_length)
            building_name = self._select_building(fold, self.target_appliance)
            activations = (
                self.activations[fold][self.target_appliance][building_name])
            activation_i = self._select_activation(activations)
            activation = activations[activation_i]
            positioned_activation, is_complete = (
                self._position_activation(
                    activation, is_target_appliance=True))
            if is_complete or self.include_incomplete_target_in_output:
                seq.target = positioned_activation
            else:
                seq.target = pd.Series(0, index=positioned_activation.index)

            # Check neighbouring activations
            mains_start = positioned_activation.index[0]
            mains_end = positioned_activation.index[-1]

            def neighbours_ok(neighbour_indicies):
                for i in neighbour_indicies:
                    activation = activations[i]
                    activation_duration = (
                        activation.index[-1] - activation.index[0])
                    neighbouring_activation_is_inside_mains_window = (
                        activation.index[0] >
                        (mains_start - activation_duration)
                        and activation.index[0] < mains_end)

                    if neighbouring_activation_is_inside_mains_window:
                        if self.allow_multiple_target_activations_in_aggregate:
                            if self.include_multiple_targets_in_output:
                                sum_target = seq.target.add(
                                    activation, fill_value=0)
                                is_complete = (
                                    sum_target.index == seq.target.index)
                                if self.allow_incomplete_target or is_complete:
                                    seq.target = sum_target[seq.target.index]
                        else:
                            return False  # need to retry
                    else:
                        return True  # everything checks out OK so far
                return True

            # Check forwards
            if not neighbours_ok(range(activation_i+1, len(activations))):
                continue
            # Check backwards
            if not neighbours_ok(range(activation_i-1, -1, -1)):
                continue

            # Get mains
            mains_for_building = self.mains[fold][activation.building]
            # load some additional data to make sure we have enough samples
            mains_end_extended = mains_end + timedelta(
                seconds=self.sample_period * 2)
            mains = mains_for_building[mains_start:mains_end_extended]
            seq.input = mains.values[:self.seq_length]
            if len(seq.input) != self.seq_length:
                continue
            return seq
        raise RuntimeError("No valid sequences found after {} retries!"
                           .format(MAX_RETRIES))

    @classmethod
    def _attrs_to_remove_for_report(cls):
        return [
            'activations', 'rng', 'mains', 'mains_good_sections',
            'sections_with_no_target', 'all_gaps']
<|endoftext|>


<commit_before>"""Tests for IRASA functions."""

import numpy as np

from neurodsp.tests.settings import FS, N_SECONDS_LONG, EXP1

from neurodsp.sim import sim_combined
from neurodsp.spectral import compute_spectrum, trim_spectrum

from neurodsp.aperiodic.irasa import *

###################################################################################################
###################################################################################################

def test_compute_irasa(tsig_comb):

    # Estimate periodic and aperiodic components with IRASA
    f_range = [1, 30]
    freqs, psd_ap, psd_pe = compute_irasa(tsig_comb, FS, f_range, noverlap=int(2*FS))

    assert len(freqs) == len(psd_ap) == len(psd_pe)

    # Compute r-squared for the full model, comparing to a standard power spectrum
    _, powers = trim_spectrum(*compute_spectrum(tsig_comb, FS, nperseg=int(4*FS)), f_range)
    r_sq = np.corrcoef(np.array([powers, psd_ap+psd_pe]))[0][1]
    assert r_sq > .95

def test_fit_irasa(tsig_comb):

    # Estimate periodic and aperiodic components with IRASA & fit aperiodic
    freqs, psd_ap, _ = compute_irasa(tsig_comb, FS, noverlap=int(2*FS))
    b0, b1 = fit_irasa(freqs, psd_ap)

    assert round(b1) == EXP1
    assert np.abs(b0 - np.log((psd_ap)[0])) < 1

def test_fit_func():

    freqs = np.arange(30)
    intercept = -2
    slope = -2

    fit = fit_func(freqs, intercept, slope)
    assert (fit == slope * freqs + intercept).all()
<commit_msg>Update test to match fix in `fit_irasa`<commit_after>"""Tests for IRASA functions."""

import numpy as np

from neurodsp.tests.settings import FS, N_SECONDS_LONG, EXP1

from neurodsp.sim import sim_combined
from neurodsp.spectral import compute_spectrum, trim_spectrum

from neurodsp.aperiodic.irasa import *

###################################################################################################
###################################################################################################

def test_compute_irasa(tsig_comb):

    # Estimate periodic and aperiodic components with IRASA
    f_range = [1, 30]
    freqs, psd_ap, psd_pe = compute_irasa(tsig_comb, FS, f_range, noverlap=int(2*FS))

    assert len(freqs) == len(psd_ap) == len(psd_pe)

    # Compute r-squared for the full model, comparing to a standard power spectrum
    _, powers = trim_spectrum(*compute_spectrum(tsig_comb, FS, nperseg=int(4*FS)), f_range)
    r_sq = np.corrcoef(np.array([powers, psd_ap+psd_pe]))[0][1]
    assert r_sq > .95

def test_fit_irasa(tsig_comb):

    # Estimate periodic and aperiodic components with IRASA & fit aperiodic
    freqs, psd_ap, _ = compute_irasa(tsig_comb, FS, noverlap=int(2*FS))
    b0, b1 = fit_irasa(freqs, psd_ap)

    assert round(b1) == EXP1
    assert np.abs(b0 - np.log10((psd_ap)[0])) < 1

def test_fit_func():

    freqs = np.arange(30)
    intercept = -2
    slope = -2

    fit = fit_func(freqs, intercept, slope)
    assert (fit == slope * freqs + intercept).all()
<|endoftext|>