├── .editorconfig ├── .gitignore ├── LICENSE.txt ├── README.txt ├── __init__.py ├── baseline.bytes ├── examples ├── console.py ├── console_recorder.py ├── example_startup.py ├── pygame_background.png └── pygame_viewer.py ├── mindwave ├── __init__.py ├── bluetooth_headset.py ├── parser.py └── pyeeg.py ├── setup.py └── test └── tests.py /.editorconfig: -------------------------------------------------------------------------------- 1 | root = true 2 | 3 | [*.py] 4 | indent_style = space 5 | indent_size = 4 6 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.swp 2 | *.pyc 3 | *Session.vim* 4 | *.o 5 | *.a 6 | *.so 7 | *.log 8 | cscope.* 9 | .*.swo 10 | tags 11 | 12 | 13 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Software License Agreement (BSD License) 2 | 3 | Copyright (c) 2014, Andreas Klostermann 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without modification, 7 | are permitted provided that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above copyright notice, this 13 | list of conditions and the following disclaimer in the documentation and/or 14 | other materials provided with the distribution. 15 | 16 | 17 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 18 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 19 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 20 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR 21 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 22 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 23 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 24 | ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 25 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 26 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 27 | -------------------------------------------------------------------------------- /README.txt: -------------------------------------------------------------------------------- 1 | Python Mindwave 2 | 3 | This is a library to interface with the Neurosky Mindwave headsets. It can 4 | parse the binary protocol spoken by Neurosky's ThinkGear AM modules. These 5 | modules are used in several headsets and board games, and development kits 6 | are available from Neurosky. 7 | 8 | This Software is licensed under the BSD License, and the authors do not 9 | represent Neurosky. 10 | 11 | This software is not intended to be used in medical diagnostics or medical 12 | treatment. 13 | 14 | 15 | Currently this software can only connect to the bluetooth version 16 | (Mindwave Mobile), because I don't posess the other variations. 17 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /baseline.bytes: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/akloster/python-mindwave/d8d8d50dddaf7dfa14ced79d25c14fe26822ff39/baseline.bytes -------------------------------------------------------------------------------- /examples/console.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | # -*- coding: utf-8 -*- 3 | 4 | from mindwave.parser import ThinkGearParser, TimeSeriesRecorder 5 | import bluetooth 6 | import time 7 | import sys 8 | import argparse 9 | 10 | 11 | from mindwave.bluetooth_headset import connect_magic, connect_bluetooth_addr 12 | from mindwave.bluetooth_headset import BluetoothError 13 | from example_startup import mindwave_startup 14 | 15 | description = """Simple Neurofeedback console application. 16 | 17 | Make sure you paired the Mindwave to your computer. You need to 18 | do that pairing for every operating system/user profile you run 19 | seperately. 20 | 21 | If you don't know the address, leave it out, and this program will 22 | figure it out, but you need to put the MindWave Mobile headset into 23 | pairing mode first. 24 | 25 | """ 26 | if __name__ == '__main__': 27 | extra_args=[dict(name='measure', type=str, nargs='?', 28 | const="attention", default="attention", 29 | help="""Measure you want feedback on. Either "meditation" 30 | or "attention\"""")] 31 | socket, args = mindwave_startup(description=description, 32 | extra_args=extra_args) 33 | 34 | if args.measure not in ["attention", "meditation"]: 35 | print("Unknown measure %s" % repr(args.measure)) 36 | sys.exit(-1) 37 | recorder = TimeSeriesRecorder() 38 | parser = ThinkGearParser(recorders=[recorder]) 39 | 40 | if args.measure== 'attention': 41 | measure_name = 'Attention' 42 | else: 43 | measure_name = 'Meditation' 44 | 45 | while 1: 46 | time.sleep(0.25) 47 | data = socket.recv(20000) 48 | parser.feed(data) 49 | v = 0 50 | if args.measure == 'attention': 51 | if len(recorder.attention)>0: 52 | v = recorder.attention[-1] 53 | if args.measure == 'meditation': 54 | if len(recorder.meditation)>0: 55 | v = recorder.meditation[-1] 56 | if v>0: 57 | print("BALABLA:",v) 58 | -------------------------------------------------------------------------------- /examples/console_recorder.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | # -*- coding: utf-8 -*- 3 | 4 | from mindwave.parser import ThinkGearParser, TimeSeriesRecorder 5 | import bluetooth 6 | import time 7 | import sys 8 | import argparse 9 | 10 | from mindwave.bluetooth_headset import connect_magic, connect_bluetooth_addr 11 | from mindwave.bluetooth_headset import BluetoothError 12 | from example_startup import mindwave_startup 13 | 14 | description = """Simple commandline application to record EEG. 15 | 16 | Make sure you paired the Mindwave to your computer. You need to 17 | do that pairing for every operating system/user profile you run 18 | seperately. 19 | 20 | If you don't know the address, leave it out, and this program will 21 | figure it out, but you need to put the MindWave Mobile headset into 22 | pairing mode first, otherwise it can't be found. 23 | 24 | """ 25 | if __name__ == '__main__': 26 | extra_args = [dict(name='filename', type=str, nargs=1, help="File to write data to (HDF5 format)."), dict(name='frequency', type=int, nargs='?', 27 | const=10, default=10, 28 | help="""Frequency of recording, in iterations per second. 29 | This doesn't affect the sampling accuracy, but rather how 30 | often the parser is translating the data from the device 31 | into Timeseries data. 32 | """)] 33 | 34 | socket, args = mindwave_startup(description=description, 35 | extra_args=extra_args) 36 | 37 | recorder = TimeSeriesRecorder(args.filename[0]) 38 | parser = ThinkGearParser(recorders=[recorder]) 39 | loop_time = 1.0 / float(args.frequency) 40 | last_message = time.time() 41 | start = last_message 42 | while 1: 43 | t = time.time() 44 | try: 45 | data = socket.recv(10000) 46 | except BluetoothError: 47 | print("BluetoothError") 48 | time.sleep(0.5) 49 | continue 50 | parser.feed(data) 51 | elapsed = time.time()-t 52 | time.sleep(max(0.01, loop_time-elapsed)) 53 | if (time.time()-last_message)>=5: 54 | print("%.2f" % (time.time()-start)) 55 | last_message = time.time() 56 | r = parser.recorders[0] 57 | print(r.meditation[-1], r.attention[-1], sum(r.raw[-100:-1])) 58 | 59 | 60 | -------------------------------------------------------------------------------- /examples/example_startup.py: -------------------------------------------------------------------------------- 1 | import time 2 | import sys 3 | import argparse 4 | import bluetooth 5 | 6 | from mindwave.bluetooth_headset import connect_magic, connect_bluetooth_addr 7 | from mindwave.bluetooth_headset import BluetoothError 8 | 9 | def mindwave_startup(description="", extra_args=[]): 10 | parser = argparse.ArgumentParser(description=description) 11 | parser.add_argument('address', type=str, nargs='?', 12 | const=None, default=None, 13 | help="""Bluetooth Address of device. Use this 14 | if you have multiple headsets nearby or you want 15 | to save a few seconds during startup.""") 16 | for params in extra_args: 17 | name = params['name'] 18 | del params['name'] 19 | parser.add_argument(name, **params) 20 | args = parser.parse_args(sys.argv[1:]) 21 | if args.address is None: 22 | socket, socket_addr = connect_magic() 23 | if socket is None: 24 | print( "No MindWave Mobile found.") 25 | sys.exit(-1) 26 | else: 27 | socket = connect_bluetooth_addr(args.address) 28 | if socket is None: 29 | print("Connection failed.") 30 | sys.exit(-1) 31 | socket_addr = args.address 32 | print( "Connected with MindWave Mobile at %s" % socket_addr) 33 | for i in range(5): 34 | try: 35 | if i>0: 36 | print("Retrying...") 37 | time.sleep(2) 38 | len(socket.recv(10)) 39 | break 40 | except BluetoothError: 41 | print("BluetoothError") 42 | if i == 5: 43 | print( "Connection failed.") 44 | sys.exit(-1) 45 | return socket, args 46 | -------------------------------------------------------------------------------- /examples/pygame_background.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/akloster/python-mindwave/d8d8d50dddaf7dfa14ced79d25c14fe26822ff39/examples/pygame_background.png -------------------------------------------------------------------------------- /examples/pygame_viewer.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | # -*- coding: utf-8 -*- 3 | import pygame, sys 4 | from numpy import * 5 | from pygame import * 6 | import scipy 7 | from mindwave.pyeeg import bin_power 8 | from mindwave.parser import ThinkGearParser, TimeSeriesRecorder 9 | from mindwave.bluetooth_headset import connect_magic, connect_bluetooth_addr 10 | from mindwave.bluetooth_headset import BluetoothError 11 | from example_startup import mindwave_startup 12 | 13 | description = """Pygame Example 14 | """ 15 | 16 | 17 | socket, args = mindwave_startup(description=description) 18 | recorder = TimeSeriesRecorder() 19 | parser = ThinkGearParser(recorders= [recorder]) 20 | 21 | def main(): 22 | pygame.init() 23 | 24 | fpsClock= pygame.time.Clock() 25 | 26 | window = pygame.display.set_mode((1280,720)) 27 | pygame.display.set_caption("Mindwave Viewer") 28 | 29 | 30 | blackColor = pygame.Color(0,0,0) 31 | redColor = pygame.Color(255,0,0) 32 | greenColor = pygame.Color(0,255,0) 33 | deltaColor = pygame.Color(100,0,0) 34 | thetaColor = pygame.Color(0,0,255) 35 | alphaColor = pygame.Color(255,0,0) 36 | betaColor = pygame.Color(0,255,00) 37 | gammaColor = pygame.Color(0,255,255) 38 | 39 | 40 | background_img = pygame.image.load("pygame_background.png") 41 | 42 | 43 | font = pygame.font.Font("freesansbold.ttf", 20) 44 | raw_eeg = True 45 | spectra = [] 46 | iteration = 0 47 | 48 | meditation_img = font.render("Meditation", False, redColor) 49 | attention_img = font.render("Attention", False, redColor) 50 | 51 | record_baseline = False 52 | quit = False 53 | while quit is False: 54 | try: 55 | data = socket.recv(10000) 56 | parser.feed(data) 57 | except BluetoothError: 58 | pass 59 | window.blit(background_img,(0,0)) 60 | if len(recorder.attention)>0: 61 | iteration+=1 62 | flen = 50 63 | if len(recorder.raw)>=500: 64 | spectrum, relative_spectrum = bin_power(recorder.raw[-512*3:], range(flen),512) 65 | spectra.append(array(relative_spectrum)) 66 | if len(spectra)>30: 67 | spectra.pop(0) 68 | 69 | spectrum = mean(array(spectra),axis=0) 70 | for i in range (flen-1): 71 | value = float(spectrum[i]*1000) 72 | if i<3: 73 | color = deltaColor 74 | elif i<8: 75 | color = thetaColor 76 | elif i<13: 77 | color = alphaColor 78 | elif i<30: 79 | color = betaColor 80 | else: 81 | color = gammaColor 82 | pygame.draw.rect(window, color, (25+i*10, 400-value, 5, value)) 83 | else: 84 | pass 85 | pygame.draw.circle(window, redColor, (800,200), int(recorder.attention[-1]/2)) 86 | pygame.draw.circle(window, greenColor, (800,200), 60//2,1) 87 | pygame.draw.circle(window, greenColor, (800,200), 100//2,1) 88 | window.blit(attention_img, (760,260)) 89 | pygame.draw.circle(window, redColor, (700,200), int(recorder.meditation[-1]/2)) 90 | pygame.draw.circle(window, greenColor, (700,200), 60//2, 1) 91 | pygame.draw.circle(window, greenColor, (700,200), 100//2, 1) 92 | 93 | window.blit(meditation_img, (600,260)) 94 | 95 | """if len(parser.current_vector)>7: 96 | m = max(p.current_vector) 97 | for i in range(7): 98 | if m == 0: 99 | value = 0 100 | else: 101 | value = p.current_vector[i] *100.0/m 102 | pygame.draw.rect(window, redColor, (600+i*30,450-value, 6,value))""" 103 | if raw_eeg: 104 | lv = 0 105 | for i,value in enumerate(recorder.raw[-1000:]): 106 | v = value/ 2.0 107 | pygame.draw.line(window, redColor, (i+25, 500-lv), (i+25, 500-v)) 108 | lv = v 109 | else: 110 | img = font.render("Not receiving any data from mindwave...", False, redColor) 111 | window.blit(img,(100,100)) 112 | pass 113 | 114 | for event in pygame.event.get(): 115 | if event.type==QUIT: 116 | quit = True 117 | if event.type==KEYDOWN: 118 | if event.key==K_ESCAPE: 119 | quit = True 120 | pygame.display.update() 121 | fpsClock.tick(12) 122 | 123 | if __name__ == '__main__': 124 | try: 125 | main() 126 | finally: 127 | pygame.quit() 128 | -------------------------------------------------------------------------------- /mindwave/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /mindwave/bluetooth_headset.py: -------------------------------------------------------------------------------- 1 | import bluetooth 2 | from bluetooth.btcommon import BluetoothError 3 | import json 4 | import time 5 | 6 | def connect_bluetooth_addr(addr): 7 | for i in range(5): 8 | if i > 0: 9 | time.sleep(1) 10 | sock = bluetooth.BluetoothSocket(bluetooth.RFCOMM) 11 | try: 12 | sock.connect((addr, 1)) 13 | sock.setblocking(False) 14 | return sock 15 | except BluetoothError: 16 | print("ERROR!") 17 | raise 18 | return None 19 | 20 | def connect_magic(): 21 | """ Tries to connect to the first MindWave Mobile it can find. 22 | If this computer hasn't connected to the headset before, you may need 23 | to make it visible to this computer's bluetooth adapter. This is done 24 | by pushing the switch on the left side of the headset to the "on/pair" 25 | position until the blinking rhythm switches. 26 | 27 | The address is then put in a file for later reference. 28 | 29 | """ 30 | return (connect_bluetooth_addr("0D:00:18:21:64:1E"),"0D:00:18:21:64:1E") 31 | # nearby_devices = bluetooth.discover_devices(lookup_names = True, duration=5) 32 | 33 | # for addr, name in nearby_devices: 34 | # print(addr,name) 35 | # if addr == '0D:00:18:21:64:1E': 36 | # print("found") 37 | # return (connect_bluetooth_addr(addr), addr) 38 | # return (None, "") 39 | -------------------------------------------------------------------------------- /mindwave/parser.py: -------------------------------------------------------------------------------- 1 | import bluetooth 2 | import struct 3 | import time 4 | import pandas as pd 5 | from datetime import datetime 6 | 7 | 8 | """ 9 | 10 | This interface library is designed to be used from very different contexts. 11 | The general idea is that the Mindwave modules in the headset (and other devices) 12 | talk a common binary protocol, which is entirely one-sided from headset to device/ 13 | computer, with one exception (explained later). The means of transport however 14 | does vary. The original MindWave headset had 2.4Ghz wireless connection, using a 15 | proprietary USB dongle/receiver. This receiver is mounted as a serial console in 16 | Linux. It also requires extra commands to connect and disconnect. 17 | 18 | The MindWave mobile uses bluetooth, which I would recommend over the 2.4Ghz version. 19 | 20 | There have been hacks with arduinos hooked up to the Thinkgear AM modules directly. 21 | 22 | Not only are the technical means of data transport different, your application needs 23 | one of several possible means of regularly reading the data. 24 | 25 | In the EuroPython 2014 talk "Brainwaves for Hackers" I demonstrated a way to do this 26 | in the IPython Notebook, and that only involved a blocking read from a bluetooth socket at 27 | certain intervals. Pygame works the same way. 28 | 29 | There are more sophisticated event loops out there, like in Kivy, Gevent or Tornado. 30 | 31 | That are the reasons why there is a parser module that can be fed a stream of bytes. 32 | You can add recorders to the parser, which take care of analyzing the parsed data. 33 | 34 | There is for example one recorder which converts the parsed data into Pandas 35 | Timeseries. But doing that dozens of times per second is too much work for weak 36 | processors, like in the Raspberry Pi, so there you would probably derive your own 37 | parser. 38 | """ 39 | def queue_to_series(a, freq="s"): 40 | t = pd.date_range(end=datetime.now(), freq=freq, periods=len(a)) 41 | return pd.TimeSeries(a, index=t) 42 | 43 | class ThinkGearParser(object): 44 | def __init__(self, recorders=None): 45 | self.recorders = [] 46 | if recorders is not None: 47 | self.recorders += recorders 48 | self.input_data = "" 49 | self.parser = self.parse() 50 | self.parser.next() 51 | 52 | def feed(self, data): 53 | for c in data: 54 | self.parser.send(ord(c)) 55 | for recorder in self.recorders: 56 | recorder.finish_chunk() 57 | self.input_data += data 58 | def dispatch_data(self, key, value): 59 | for recorder in self.recorders: 60 | recorder.dispatch_data(key, value) 61 | 62 | def parse(self): 63 | """ 64 | This generator parses one byte at a time. 65 | """ 66 | i = 1 67 | times = [] 68 | while 1: 69 | byte = yield 70 | if byte== 0xaa: 71 | byte = yield # This byte should be "\aa" too 72 | if byte== 0xaa: 73 | # packet synced by 0xaa 0xaa 74 | packet_length = yield 75 | packet_code = yield 76 | if packet_code == 0xd4: 77 | # standing by 78 | self.state = "standby" 79 | elif packet_code == 0xd0: 80 | self.state = "connected" 81 | elif packet_code == 0xd2: 82 | data_len = yield 83 | headset_id = yield 84 | headset_id += yield 85 | self.dongle_state = "disconnected" 86 | else: 87 | self.sending_data = True 88 | left = packet_length - 2 89 | while left>0: 90 | if packet_code ==0x80: # raw value 91 | row_length = yield 92 | a = yield 93 | b = yield 94 | value = struct.unpack("0: 104 | v = struct.unpack("b",chr(a))[0] 105 | if 0 < v <= 100: 106 | self.dispatch_data("attention", v) 107 | left-=1 108 | elif packet_code == 0x05: # Meditation (eSense) 109 | a = yield 110 | if a>0: 111 | v = struct.unpack("b",chr(a))[0] 112 | if 0 < v <= 100: 113 | self.dispatch_data("meditation", v) 114 | left-=1 115 | 116 | 117 | elif packet_code == 0x16: # Blink Strength 118 | self.current_blink_strength = yield 119 | 120 | left-=1 121 | elif packet_code == 0x83: 122 | vlength = yield 123 | self.current_vector = [] 124 | for row in range(8): 125 | a = yield 126 | b = yield 127 | c = yield 128 | value = a*255*255+b*255+c 129 | left -= vlength 130 | self.dispatch_data("bands", self.current_vector) 131 | packet_code = yield 132 | else: 133 | pass # sync failed 134 | else: 135 | pass # sync failed 136 | 137 | class TimeSeriesRecorder: 138 | def __init__(self, file_name=None): 139 | self.meditation = pd.TimeSeries() 140 | self.attention = pd.TimeSeries() 141 | self.raw = pd.TimeSeries() 142 | self.blink = pd.TimeSeries() 143 | self.poor_signal = pd.TimeSeries() 144 | self.attention_queue = [] 145 | self.meditation_queue = [] 146 | self.poor_signal_queue = [] 147 | self.blink_queue = [] 148 | self.raw_queue = [] 149 | if file_name is not None: 150 | self.store = pd.HDFStore(file_name) 151 | else: 152 | self.store = None 153 | 154 | def dispatch_data(self, key, value): 155 | if key == "attention": 156 | self.attention_queue.append(value) 157 | # Blink and "poor signal" is only sent when a blink or poor signal is detected 158 | # So fake continuous signal as zeros. 159 | 160 | self.blink_queue.append(0) 161 | self.poor_signal_queue.append(0) 162 | 163 | elif key == "meditation": 164 | self.meditation_queue.append(value) 165 | elif key == "raw": 166 | self.raw_queue.append(value) 167 | elif key == "blink": 168 | self.blink_queue.append(value) 169 | if len(self.blink_queue)>0: 170 | self.blink_queue[-1] = self.current_blink_strength 171 | 172 | elif key == "poor_signal": 173 | if len(self.poor_signal_queue)>0: 174 | self.poor_signal_queue[-1] = a 175 | 176 | 177 | def record_meditation(self, attention): 178 | self.meditation_queue.append() 179 | 180 | def record_blink(self, attention): 181 | self.blink_queue.append() 182 | 183 | def finish_chunk(self): 184 | """ called periodically to update the timeseries """ 185 | self.meditation = pd.concat([self.meditation, queue_to_series(self.meditation_queue, freq="s")]) 186 | 187 | self.attention = pd.concat([self.attention, queue_to_series(self.attention_queue, freq="s")]) 188 | self.blink = pd.concat([self.blink, queue_to_series(self.blink_queue, freq="s")]) 189 | self.raw = pd.concat([self.raw, queue_to_series(self.raw_queue, freq="1953U")]) 190 | self.poor_signal = pd.concat([self.poor_signal, queue_to_series(self.poor_signal_queue)]) 191 | 192 | self.attention_queue = [] 193 | self.meditation_queue = [] 194 | self.poor_signal_queue = [] 195 | self.blink_queue = [] 196 | self.raw_queue = [] 197 | if self.store is not None: 198 | self.store['attention'] = self.attention 199 | self.store['meditation'] = self.meditation 200 | self.store['raw'] = self.raw 201 | 202 | -------------------------------------------------------------------------------- /mindwave/pyeeg.py: -------------------------------------------------------------------------------- 1 | """Copyleft 2010 Forrest Sheng Bao http://fsbao.net 2 | 3 | PyEEG, a Python module to extract EEG features, v 0.02_r2 4 | 5 | Project homepage: http://pyeeg.org 6 | 7 | **Data structure** 8 | 9 | PyEEG only uses standard Python and numpy data structures, 10 | so you need to import numpy before using it. 11 | For numpy, please visit http://numpy.scipy.org 12 | 13 | **Naming convention** 14 | 15 | I follow "Style Guide for Python Code" to code my program 16 | http://www.python.org/dev/peps/pep-0008/ 17 | 18 | Constants: UPPER_CASE_WITH_UNDERSCORES, e.g., SAMPLING_RATE, LENGTH_SIGNAL. 19 | 20 | Function names: lower_case_with_underscores, e.g., spectrum_entropy. 21 | 22 | Variables (global and local): CapitalizedWords or CapWords, e.g., Power. 23 | 24 | If a variable name consists of one letter, I may use lower case, e.g., x, y. 25 | 26 | Functions listed alphabetically 27 | -------------------------------------------------- 28 | 29 | """ 30 | 31 | from numpy.fft import fft 32 | from numpy import zeros, floor, log10, log, mean, array, sqrt, vstack, cumsum, \ 33 | ones, log2, std 34 | from numpy.linalg import svd, lstsq 35 | import time 36 | 37 | ######################## Functions contributed by Xin Liu ################# 38 | 39 | def hurst(X): 40 | """ Compute the Hurst exponent of X. If the output H=0.5,the behavior 41 | of the time-series is similar to random walk. If H<0.5, the time-series 42 | cover less "distance" than a random walk, vice verse. 43 | 44 | Parameters 45 | ---------- 46 | 47 | X 48 | 49 | list 50 | 51 | a time series 52 | 53 | Returns 54 | ------- 55 | H 56 | 57 | float 58 | 59 | Hurst exponent 60 | 61 | Examples 62 | -------- 63 | 64 | >>> import pyeeg 65 | >>> from numpy.random import randn 66 | >>> a = randn(4096) 67 | >>> pyeeg.hurst(a) 68 | >>> 0.5057444 69 | 70 | """ 71 | 72 | N = len(X) 73 | 74 | T = array([float(i) for i in range(1,N+1)]) 75 | Y = cumsum(X) 76 | Ave_T = Y/T 77 | 78 | S_T = zeros((N)) 79 | R_T = zeros((N)) 80 | for i in range(N): 81 | S_T[i] = std(X[:i+1]) 82 | X_T = Y - T * Ave_T[i] 83 | R_T[i] = max(X_T[:i + 1]) - min(X_T[:i + 1]) 84 | 85 | R_S = R_T / S_T 86 | R_S = log(R_S) 87 | n = log(T).reshape(N, 1) 88 | H = lstsq(n[1:], R_S[1:])[0] 89 | return H[0] 90 | 91 | 92 | ######################## Begin function definitions ####################### 93 | 94 | def embed_seq(X,Tau,D): 95 | """Build a set of embedding sequences from given time series X with lag Tau 96 | and embedding dimension DE. Let X = [x(1), x(2), ... , x(N)], then for each 97 | i such that 1 < i < N - (D - 1) * Tau, we build an embedding sequence, 98 | Y(i) = [x(i), x(i + Tau), ... , x(i + (D - 1) * Tau)]. All embedding 99 | sequence are placed in a matrix Y. 100 | 101 | Parameters 102 | ---------- 103 | 104 | X 105 | list 106 | 107 | a time series 108 | 109 | Tau 110 | integer 111 | 112 | the lag or delay when building embedding sequence 113 | 114 | D 115 | integer 116 | 117 | the embedding dimension 118 | 119 | Returns 120 | ------- 121 | 122 | Y 123 | 2-D list 124 | 125 | embedding matrix built 126 | 127 | Examples 128 | --------------- 129 | >>> import pyeeg 130 | >>> a=range(0,9) 131 | >>> pyeeg.embed_seq(a,1,4) 132 | array([[ 0., 1., 2., 3.], 133 | [ 1., 2., 3., 4.], 134 | [ 2., 3., 4., 5.], 135 | [ 3., 4., 5., 6.], 136 | [ 4., 5., 6., 7.], 137 | [ 5., 6., 7., 8.]]) 138 | >>> pyeeg.embed_seq(a,2,3) 139 | array([[ 0., 2., 4.], 140 | [ 1., 3., 5.], 141 | [ 2., 4., 6.], 142 | [ 3., 5., 7.], 143 | [ 4., 6., 8.]]) 144 | >>> pyeeg.embed_seq(a,4,1) 145 | array([[ 0.], 146 | [ 1.], 147 | [ 2.], 148 | [ 3.], 149 | [ 4.], 150 | [ 5.], 151 | [ 6.], 152 | [ 7.], 153 | [ 8.]]) 154 | 155 | 156 | 157 | """ 158 | N =len(X) 159 | 160 | if D * Tau > N: 161 | print "Cannot build such a matrix, because D * Tau > N" 162 | exit() 163 | 164 | if Tau<1: 165 | print "Tau has to be at least 1" 166 | exit() 167 | 168 | Y=zeros((N - (D - 1) * Tau, D)) 169 | for i in range(0, N - (D - 1) * Tau): 170 | for j in range(0, D): 171 | Y[i][j] = X[i + j * Tau] 172 | return Y 173 | 174 | def in_range(Template, Scroll, Distance): 175 | """Determines whether one vector is the the range of another vector. 176 | 177 | The two vectors should have equal length. 178 | 179 | Parameters 180 | ----------------- 181 | Template 182 | list 183 | The template vector, one of two vectors being compared 184 | 185 | Scroll 186 | list 187 | The scroll vector, one of the two vectors being compared 188 | 189 | D 190 | float 191 | Two vectors match if their distance is less than D 192 | 193 | Bit 194 | 195 | 196 | Notes 197 | ------- 198 | The distance between two vectors can be defined as Euclidean distance 199 | according to some publications. 200 | 201 | The two vector should of equal length 202 | 203 | """ 204 | 205 | for i in range(0, len(Template)): 206 | if abs(Template[i] - Scroll[i]) > Distance: 207 | return False 208 | return True 209 | """ Desperate code, but do not delete 210 | def bit_in_range(Index): 211 | if abs(Scroll[Index] - Template[Bit]) <= Distance : 212 | print "Bit=", Bit, "Scroll[Index]", Scroll[Index], "Template[Bit]",\ 213 | Template[Bit], "abs(Scroll[Index] - Template[Bit])",\ 214 | abs(Scroll[Index] - Template[Bit]) 215 | return Index + 1 # move 216 | 217 | Match_No_Tail = range(0, len(Scroll) - 1) # except the last one 218 | # print Match_No_Tail 219 | 220 | # first compare Template[:-2] and Scroll[:-2] 221 | 222 | for Bit in range(0, len(Template) - 1): # every bit of Template is in range of Scroll 223 | Match_No_Tail = filter(bit_in_range, Match_No_Tail) 224 | print Match_No_Tail 225 | 226 | # second and last, check whether Template[-1] is in range of Scroll and 227 | # Scroll[-1] in range of Template 228 | 229 | # 2.1 Check whether Template[-1] is in the range of Scroll 230 | Bit = - 1 231 | Match_All = filter(bit_in_range, Match_No_Tail) 232 | 233 | # 2.2 Check whether Scroll[-1] is in the range of Template 234 | # I just write a loop for this. 235 | for i in Match_All: 236 | if abs(Scroll[-1] - Template[i] ) <= Distance: 237 | Match_All.remove(i) 238 | 239 | 240 | return len(Match_All), len(Match_No_Tail) 241 | """ 242 | 243 | def bin_power(X,Band,Fs): 244 | """Compute power in each frequency bin specified by Band from FFT result of 245 | X. By default, X is a real signal. 246 | 247 | Note 248 | ----- 249 | A real signal can be synthesized, thus not real. 250 | 251 | Parameters 252 | ----------- 253 | 254 | Band 255 | list 256 | 257 | boundary frequencies (in Hz) of bins. They can be unequal bins, e.g. 258 | [0.5,4,7,12,30] which are delta, theta, alpha and beta respectively. 259 | You can also use range() function of Python to generate equal bins and 260 | pass the generated list to this function. 261 | 262 | Each element of Band is a physical frequency and shall not exceed the 263 | Nyquist frequency, i.e., half of sampling frequency. 264 | 265 | X 266 | list 267 | 268 | a 1-D real time series. 269 | 270 | Fs 271 | integer 272 | 273 | the sampling rate in physical frequency 274 | 275 | Returns 276 | ------- 277 | 278 | Power 279 | list 280 | 281 | spectral power in each frequency bin. 282 | 283 | Power_ratio 284 | list 285 | 286 | spectral power in each frequency bin normalized by total power in ALL 287 | frequency bins. 288 | 289 | """ 290 | 291 | C = fft(X) 292 | C = abs(C) 293 | Power =zeros(len(Band)-1); 294 | for Freq_Index in range(0,len(Band)-1): 295 | Freq = float(Band[Freq_Index]) ## Xin Liu 296 | Next_Freq = float(Band[Freq_Index+1]) 297 | Power[Freq_Index] = sum(C[floor(Freq/Fs*len(X)):floor(Next_Freq/Fs*len(X))]) 298 | Power_Ratio = Power/sum(Power) 299 | return Power, Power_Ratio 300 | 301 | def first_order_diff(X): 302 | """ Compute the first order difference of a time series. 303 | 304 | For a time series X = [x(1), x(2), ... , x(N)], its first order 305 | difference is: 306 | Y = [x(2) - x(1) , x(3) - x(2), ..., x(N) - x(N-1)] 307 | 308 | """ 309 | D=[] 310 | 311 | for i in range(1,len(X)): 312 | D.append(X[i]-X[i-1]) 313 | 314 | return D 315 | 316 | def pfd(X, D=None): 317 | """Compute Petrosian Fractal Dimension of a time series from either two 318 | cases below: 319 | 1. X, the time series of type list (default) 320 | 2. D, the first order differential sequence of X (if D is provided, 321 | recommended to speed up) 322 | 323 | In case 1, D is computed by first_order_diff(X) function of pyeeg 324 | 325 | To speed up, it is recommended to compute D before calling this function 326 | because D may also be used by other functions whereas computing it here 327 | again will slow down. 328 | """ 329 | if D is None: ## Xin Liu 330 | D = first_order_diff(X) 331 | N_delta= 0; #number of sign changes in derivative of the signal 332 | for i in range(1,len(D)): 333 | if D[i]*D[i-1]<0: 334 | N_delta += 1 335 | n = len(X) 336 | return log10(n)/(log10(n)+log10(n/n+0.4*N_delta)) 337 | 338 | 339 | def hfd(X, Kmax): 340 | """ Compute Hjorth Fractal Dimension of a time series X, kmax 341 | is an HFD parameter 342 | """ 343 | L = []; 344 | x = [] 345 | N = len(X) 346 | for k in range(1,Kmax): 347 | Lk = [] 348 | for m in range(0,k): 349 | Lmk = 0 350 | for i in range(1,int(floor((N-m)/k))): 351 | Lmk += abs(X[m+i*k] - X[m+i*k-k]) 352 | Lmk = Lmk*(N - 1)/floor((N - m) / float(k)) / k 353 | Lk.append(Lmk) 354 | L.append(log(mean(Lk))) 355 | x.append([log(float(1) / k), 1]) 356 | 357 | (p, r1, r2, s)=lstsq(x, L) 358 | return p[0] 359 | 360 | def hjorth(X, D = None): 361 | """ Compute Hjorth mobility and complexity of a time series from either two 362 | cases below: 363 | 1. X, the time series of type list (default) 364 | 2. D, a first order differential sequence of X (if D is provided, 365 | recommended to speed up) 366 | 367 | In case 1, D is computed by first_order_diff(X) function of pyeeg 368 | 369 | Notes 370 | ----- 371 | To speed up, it is recommended to compute D before calling this function 372 | because D may also be used by other functions whereas computing it here 373 | again will slow down. 374 | 375 | Parameters 376 | ---------- 377 | 378 | X 379 | list 380 | 381 | a time series 382 | 383 | D 384 | list 385 | 386 | first order differential sequence of a time series 387 | 388 | Returns 389 | ------- 390 | 391 | As indicated in return line 392 | 393 | Hjorth mobility and complexity 394 | 395 | """ 396 | 397 | if D is None: 398 | D = first_order_diff(X) 399 | 400 | D.insert(0, X[0]) # pad the first difference 401 | D = array(D) 402 | 403 | n = len(X) 404 | 405 | M2 = float(sum(D ** 2)) / n 406 | TP = sum(array(X) ** 2) 407 | M4 = 0; 408 | for i in range(1, len(D)): 409 | M4 += (D[i] - D[i - 1]) ** 2 410 | M4 = M4 / n 411 | 412 | return sqrt(M2 / TP), sqrt(float(M4) * TP / M2 / M2) # Hjorth Mobility and Complexity 413 | 414 | def spectral_entropy(X, Band, Fs, Power_Ratio = None): 415 | """Compute spectral entropy of a time series from either two cases below: 416 | 1. X, the time series (default) 417 | 2. Power_Ratio, a list of normalized signal power in a set of frequency 418 | bins defined in Band (if Power_Ratio is provided, recommended to speed up) 419 | 420 | In case 1, Power_Ratio is computed by bin_power() function. 421 | 422 | Notes 423 | ----- 424 | To speed up, it is recommended to compute Power_Ratio before calling this 425 | function because it may also be used by other functions whereas computing 426 | it here again will slow down. 427 | 428 | Parameters 429 | ---------- 430 | 431 | Band 432 | list 433 | 434 | boundary frequencies (in Hz) of bins. They can be unequal bins, e.g. 435 | [0.5,4,7,12,30] which are delta, theta, alpha and beta respectively. 436 | You can also use range() function of Python to generate equal bins and 437 | pass the generated list to this function. 438 | 439 | Each element of Band is a physical frequency and shall not exceed the 440 | Nyquist frequency, i.e., half of sampling frequency. 441 | 442 | X 443 | list 444 | 445 | a 1-D real time series. 446 | 447 | Fs 448 | integer 449 | 450 | the sampling rate in physical frequency 451 | 452 | Returns 453 | ------- 454 | 455 | As indicated in return line 456 | 457 | See Also 458 | -------- 459 | bin_power: pyeeg function that computes spectral power in frequency bins 460 | 461 | """ 462 | 463 | if Power_Ratio is None: 464 | Power, Power_Ratio = bin_power(X, Band, Fs) 465 | 466 | Spectral_Entropy = 0 467 | for i in range(0, len(Power_Ratio) - 1): 468 | Spectral_Entropy += Power_Ratio[i] * log(Power_Ratio[i]) 469 | Spectral_Entropy /= log(len(Power_Ratio)) # to save time, minus one is omitted 470 | return -1 * Spectral_Entropy 471 | 472 | def svd_entropy(X, Tau, DE, W = None): 473 | """Compute SVD Entropy from either two cases below: 474 | 1. a time series X, with lag tau and embedding dimension dE (default) 475 | 2. a list, W, of normalized singular values of a matrix (if W is provided, 476 | recommend to speed up.) 477 | 478 | If W is None, the function will do as follows to prepare singular spectrum: 479 | 480 | First, computer an embedding matrix from X, Tau and DE using pyeeg 481 | function embed_seq(): 482 | M = embed_seq(X, Tau, DE) 483 | 484 | Second, use scipy.linalg function svd to decompose the embedding matrix 485 | M and obtain a list of singular values: 486 | W = svd(M, compute_uv=0) 487 | 488 | At last, normalize W: 489 | W /= sum(W) 490 | 491 | Notes 492 | ------------- 493 | 494 | To speed up, it is recommended to compute W before calling this function 495 | because W may also be used by other functions whereas computing it here 496 | again will slow down. 497 | """ 498 | 499 | if W is None: 500 | Y = EmbedSeq(X, tau, dE) 501 | W = svd(Y, compute_uv = 0) 502 | W /= sum(W) # normalize singular values 503 | 504 | return -1*sum(W * log(W)) 505 | 506 | def fisher_info(X, Tau, DE, W = None): 507 | """ Compute Fisher information of a time series from either two cases below: 508 | 1. X, a time series, with lag Tau and embedding dimension DE (default) 509 | 2. W, a list of normalized singular values, i.e., singular spectrum (if W is 510 | provided, recommended to speed up.) 511 | 512 | If W is None, the function will do as follows to prepare singular spectrum: 513 | 514 | First, computer an embedding matrix from X, Tau and DE using pyeeg 515 | function embed_seq(): 516 | M = embed_seq(X, Tau, DE) 517 | 518 | Second, use scipy.linalg function svd to decompose the embedding matrix 519 | M and obtain a list of singular values: 520 | W = svd(M, compute_uv=0) 521 | 522 | At last, normalize W: 523 | W /= sum(W) 524 | 525 | Parameters 526 | ---------- 527 | 528 | X 529 | list 530 | 531 | a time series. X will be used to build embedding matrix and compute 532 | singular values if W or M is not provided. 533 | Tau 534 | integer 535 | 536 | the lag or delay when building a embedding sequence. Tau will be used 537 | to build embedding matrix and compute singular values if W or M is not 538 | provided. 539 | DE 540 | integer 541 | 542 | the embedding dimension to build an embedding matrix from a given 543 | series. DE will be used to build embedding matrix and compute 544 | singular values if W or M is not provided. 545 | W 546 | list or array 547 | 548 | the set of singular values, i.e., the singular spectrum 549 | 550 | Returns 551 | ------- 552 | 553 | FI 554 | integer 555 | 556 | Fisher information 557 | 558 | Notes 559 | ----- 560 | To speed up, it is recommended to compute W before calling this function 561 | because W may also be used by other functions whereas computing it here 562 | again will slow down. 563 | 564 | See Also 565 | -------- 566 | embed_seq : embed a time series into a matrix 567 | """ 568 | 569 | if W is None: 570 | M = embed_seq(X, Tau, DE) 571 | W = svd(M, compute_uv = 0) 572 | W /= sum(W) 573 | 574 | FI = 0 575 | for i in range(0, len(W) - 1): # from 1 to M 576 | FI += ((W[i +1] - W[i]) ** 2) / (W[i]) 577 | 578 | return FI 579 | 580 | def ap_entropy(X, M, R): 581 | """Computer approximate entropy (ApEN) of series X, specified by M and R. 582 | 583 | Suppose given time series is X = [x(1), x(2), ... , x(N)]. We first build 584 | embedding matrix Em, of dimension (N-M+1)-by-M, such that the i-th row of Em 585 | is x(i),x(i+1), ... , x(i+M-1). Hence, the embedding lag and dimension are 586 | 1 and M-1 respectively. Such a matrix can be built by calling pyeeg function 587 | as Em = embed_seq(X, 1, M). Then we build matrix Emp, whose only 588 | difference with Em is that the length of each embedding sequence is M + 1 589 | 590 | Denote the i-th and j-th row of Em as Em[i] and Em[j]. Their k-th elements 591 | are Em[i][k] and Em[j][k] respectively. The distance between Em[i] and Em[j] 592 | is defined as 1) the maximum difference of their corresponding scalar 593 | components, thus, max(Em[i]-Em[j]), or 2) Euclidean distance. We say two 1-D 594 | vectors Em[i] and Em[j] *match* in *tolerance* R, if the distance between them 595 | is no greater than R, thus, max(Em[i]-Em[j]) <= R. Mostly, the value of R is 596 | defined as 20% - 30% of standard deviation of X. 597 | 598 | Pick Em[i] as a template, for all j such that 0 < j < N - M + 1, we can 599 | check whether Em[j] matches with Em[i]. Denote the number of Em[j], 600 | which is in the range of Em[i], as k[i], which is the i-th element of the 601 | vector k. The probability that a random row in Em matches Em[i] is 602 | \simga_1^{N-M+1} k[i] / (N - M + 1), thus sum(k)/ (N - M + 1), 603 | denoted as Cm[i]. 604 | 605 | We repeat the same process on Emp and obtained Cmp[i], but here 0>> import pyeeg 805 | >>> from numpy.random import randn 806 | >>> print pyeeg.dfa(randn(4096)) 807 | 0.490035110345 808 | 809 | Reference 810 | --------- 811 | Peng C-K, Havlin S, Stanley HE, Goldberger AL. Quantification of scaling 812 | exponents and crossover phenomena in nonstationary heartbeat time series. 813 | _Chaos_ 1995;5:82-87 814 | 815 | Notes 816 | ----- 817 | 818 | This value depends on the box sizes very much. When the input is a white 819 | noise, this value should be 0.5. But, some choices on box sizes can lead to 820 | the value lower or higher than 0.5, e.g. 0.38 or 0.58. 821 | 822 | Based on many test, I set the box sizes from 1/5 of signal length to one 823 | (x-5)-th of the signal length, where x is the nearest power of 2 from the 824 | length of the signal, i.e., 1/16, 1/32, 1/64, 1/128, ... 825 | 826 | You may generate a list of box sizes and pass in such a list as a parameter. 827 | 828 | """ 829 | 830 | X = array(X) 831 | 832 | if Ave is None: 833 | Ave = mean(X) 834 | 835 | Y = cumsum(X) 836 | Y -= Ave 837 | 838 | if L is None: 839 | L = floor(len(X)*1/(2**array(range(4,int(log2(len(X)))-4)))) 840 | 841 | F = zeros(len(L)) # F(n) of different given box length n 842 | 843 | for i in range(0,len(L)): 844 | n = int(L[i]) # for each box length L[i] 845 | if n==0: 846 | print "time series is too short while the box length is too big" 847 | print "abort" 848 | exit() 849 | for j in range(0,len(X),n): # for each box 850 | if j+n < len(X): 851 | c = range(j,j+n) 852 | c = vstack([c, ones(n)]).T # coordinates of time in the box 853 | y = Y[j:j+n] # the value of data in the box 854 | F[i] += lstsq(c,y)[1] # add residue in this box 855 | F[i] /= ((len(X)/n)*n) 856 | F = sqrt(F) 857 | 858 | Alpha = lstsq(vstack([log(L), ones(len(L))]).T,log(F))[0][0] 859 | 860 | return Alpha 861 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup 2 | 3 | setup( 4 | name='mindwave', 5 | version='0.2dev', 6 | packages=['mindwave',], 7 | license='BSD License', 8 | long_description=open('README.txt').read(), 9 | ) 10 | -------------------------------------------------------------------------------- /test/tests.py: -------------------------------------------------------------------------------- 1 | from mindwave.parser import ThinkGearParser, TimeSeriesRecorder 2 | 3 | import unittest 4 | 5 | class ParserTests(unittest.TestCase): 6 | 7 | def testParser(self): 8 | s = file("baseline.bytes").read() 9 | ts_recorder = TimeSeriesRecorder() 10 | parser = ThinkGearParser(recorders=[ts_recorder]) 11 | parser.feed(s) 12 | 13 | self.assertTrue(len(ts_recorder.attention) == len(ts_recorder.meditation) == len(ts_recorder.blink) == len(ts_recorder.poor_signal)) 14 | 15 | 16 | 17 | def main(): 18 | unittest.main() 19 | 20 | if __name__ == '__main__': 21 | main() 22 | --------------------------------------------------------------------------------