From f4882a3c3bb189776ea022b730c9c40201aacb04 Mon Sep 17 00:00:00 2001
From: vigji <luigi.petrucco@gmail.com>
Date: Wed, 10 Feb 2021 12:58:27 +0100
Subject: [PATCH] Arduino control

Arduino control for odors delivery
---
 stytra/examples/arduino_exp.py        |  22 ++++++
 stytra/hardware/external_pyfirmata.py | 110 ++++++++++++++++++++++++++
 stytra/stimulation/stimuli/arduino.py |  27 +++++++
 3 files changed, 159 insertions(+)
 create mode 100644 stytra/examples/arduino_exp.py
 create mode 100644 stytra/hardware/external_pyfirmata.py
 create mode 100644 stytra/stimulation/stimuli/arduino.py

diff --git a/stytra/examples/arduino_exp.py b/stytra/examples/arduino_exp.py
new file mode 100644
index 00000000..18b0602f
--- /dev/null
+++ b/stytra/examples/arduino_exp.py
@@ -0,0 +1,22 @@
+from stytra import Stytra, Protocol
+from stytra.stimulation.stimuli.set_arduino_pin import WriteArduinoPin
+
+LAYOUT = (dict(pin=5, mode="pwm", ad="d"),
+              dict(pin=11, mode="pwm", ad="d"))
+
+
+class MotorProtocol(Protocol):
+    name = 'motor_protocol'
+
+    def get_stim_sequence(self):
+        stimuli = [
+            WriteArduinoPin(port='COM3', layout=LAYOUT, pin=5, value=0.8, duration=4.0),
+            WriteArduinoPin(port='COM3', layout=LAYOUT, pin=5, value=0)
+        ]
+        return stimuli
+
+
+if __name__ == "__main__":
+
+    st = WriteArduinoPin(port='COM3', layout=LAYOUT, pin=5, value=0.8, duration=4.0)
+    #Stytra(protocol=MotorProtocol())
\ No newline at end of file
diff --git a/stytra/hardware/external_pyfirmata.py b/stytra/hardware/external_pyfirmata.py
new file mode 100644
index 00000000..73d0272c
--- /dev/null
+++ b/stytra/hardware/external_pyfirmata.py
@@ -0,0 +1,110 @@
+try:
+    import pyfirmata as pyfi
+except ImportError:
+    print('Pyfirmata is not installed')
+
+# direct mode to call pins with pyfirmata
+from time import sleep
+
+MODE_DICT = dict(input=pyfi.INPUT, output=pyfi.OUTPUT, pwm=pyfi.PWM, servo=pyfi.SERVO)
+AD_ATTR_DICT = dict(a="analog", d="digital")
+
+
+class PyfirmataConnection:
+    '''Using StandardFirmata and pyfirmata to control an Arduino board with Python'''
+
+    def __init__(self, com_port, layout, iterator=True):
+        '''If using analog ports is handy to start an iterator thread'''
+
+        self.board = pyfi.Arduino(com_port)
+        print('Connection successfully established')
+        sleep(1)
+
+    # Convert configuration list of dictionaries in a dictionary of configurations:
+        self.layout = dict()
+        for pin_conf in layout:
+            pin_conf["mode"] = MODE_DICT[pin_conf["mode"]]
+            self.layout[pin_conf.pop("pin")] = pin_conf
+
+        if iterator:
+            self.it = pyfi.util.Iterator(self.board)
+            self.it.start()
+
+        self.initialize()
+
+    def initialize(self):
+        """
+        """
+        for pin_n, pin_conf in self.layout.items():
+            getattr(self.board, AD_ATTR_DICT[pin_conf["ad"]])[pin_n].mode = pin_conf["mode"]
+
+    def read(self, pin_n):
+        """Access and read from a pin.
+
+        Parameters
+        ----------
+        pin_n : int
+            Integer indicating pin to be read
+
+        Returns
+        -------
+        float or int
+
+        """
+
+        pin_conf = self.layout[pin_n]
+        value = getattr(self.board, AD_ATTR_DICT[pin_conf["ad"]])[pin_n].read()
+
+        return value
+
+    def read_all(self):
+        """Read all pins that are set as inputs
+
+        Returns
+        -------
+        list of floats or ints
+
+        """
+        return [self.read(pin_n) for pin_n, pin_conf in self.layout.items() if pin_conf["mode"]=="input"]
+
+    def write(self, pin_n, value):
+        """
+
+        Parameters
+        ----------
+        pin_n:
+            Integer indicating pin to be written
+
+        Returns
+        -------
+
+        """
+        pin_conf = self.layout[pin_n]
+        sel_pins = getattr(self.board, AD_ATTR_DICT[pin_conf["ad"]])[pin_n]
+        sel_pins.write(value)
+
+    def write_multiple(self, values=None):
+        for pin_n, value in values.items():
+            self.write(pin_n, value)
+        return True
+
+    def close(self):
+        """Close connection"""
+
+        self.board.exit()
+        return print('Communication successfully interrupted')
+
+
+if __name__ == "__main__":
+    LAYOUT = (dict(pin=5, mode="pwm", ad="d"),
+              dict(pin=11, mode="pwm", ad="d"))
+
+    write_multiple = {5: 0,
+                      11: 0}
+
+    try_pumps = PyfirmataConnection(com_port='COM3', layout=LAYOUT)
+    #try_pumps.write(5, 0)
+
+    sleep(5)
+    try_pumps.write_multiple(write_multiple)
+    print("sending pulse")
\ No newline at end of file
diff --git a/stytra/stimulation/stimuli/arduino.py b/stytra/stimulation/stimuli/arduino.py
new file mode 100644
index 00000000..58d01ee8
--- /dev/null
+++ b/stytra/stimulation/stimuli/arduino.py
@@ -0,0 +1,27 @@
+from stytra.stimulation.stimuli import Stimulus
+from stytra.hardware.external_pyfirmata import PyfirmataConnection
+
+
+class WriteArduinoPin(Stimulus):
+    """Simple class to write a value on an arduino pin. Mostly a simple example to implement
+    your own fancy arduino classes.
+
+    Parameters
+    ----------
+    pin : int
+        Pin number.
+    value : float or int
+        Value to be set on the pin.
+    """
+
+    name = "set_arduino_pin"
+
+    def __init__(self, pin_values_dict, *args, **kwargs):
+        self.pin_values = pin_values_dict
+        super().__init__(*args, **kwargs)
+
+    def start(self):
+        super().start()
+        self._experiment.arduino_board.write_multiple(self.pin_values)
+
+