Add olfaction tutorial

.github/workflows/tests.yaml

@@ -28,7 +28,7 @@ jobs:
       - name: Lint with ruff
         run: |
           # stop the build if there are Python syntax errors or undefined names
-          ruff --format=github --select=E9,F63,F7,F82 --target-version=py38 .
+          ruff --output-format=github --select=E9,F63,F7,F82 --target-version=py38 .
       - name: Test with pytest
         run: |
           python -m pytest

doc/source/index.rst

@@ -61,7 +61,7 @@ If you use FlyGym or NeuroMechFly in your research, please cite the following tw
       title = {NeuroMechFly 2.0, a framework for simulating embodied sensorimotor control in adult Drosophila},
       year = {2023},
       doi = {10.1101/2023.09.18.556649},
-      URL = {https://www.biorxiv.org/content/early/2023/09/18/2023.09.18.556649},
+      URL = {https://www.biorxiv.org/content/10.1101/2023.09.18.556649},
       journal = {bioRxiv}
    }
 

doc/source/tutorials/olfaction.rst

@@ -1,4 +1,298 @@
 Olfaction
 =========
 
-Content forthcoming. Please refer to the API reference and our `paper  <https://www.biorxiv.org/content/10.1101/2023.09.18.556649>`_ for now.
+**Summary:** In this tutorial, we will implement a simple controller for
+odor-guided taxis.
+
+In the `previous
+tutorial <https://neuromechfly.org/tutorials/vision.html>`__, we have
+covered how one might simulate the visual experience of the fly using
+NeuroMechFly. In addition to vision, we also made it possible for our
+model to detect odors emmited by objects in the simulation environment.
+The olfactory system in *Drosophila* consists of specialized olfactory
+sensory neurons (OSNs) located in the antennae and maxillary palps.
+These detect specific odorant molecules and convey this information to
+the brain’s antennal lobe, where their signals are further processed.
+This is shown in the figure below (left, source: `Martin et al,
+2013 <https://doi.org/10.1002/ar.22747>`__) We emulated peripheral
+olfaction by attaching virtual odor sensors to the antennae and
+maxillary palps of our biomechanical model, as shown in the figure
+(right). The user has the option of configuring additional sensors at
+more precise locations on these olfactory organs. These virtual sensors
+can detect odor intensities across a multi-dimensional space that can be
+thought of as representing, for example, the concentrations of
+monomolecular chemicals sensed by OSNs in the antennae, or the
+intensities of composite odors co-activating numerous projection neurons
+in the antennal lobe.
+
+.. figure :: https://github.com/NeLy-EPFL/_media/blob/main/flygym/olfaction.png?raw=true
+   :width: 600
+
+
+The odor arena
+--------------
+
+To demonstrate odor sensing, let’s create an environment with an
+attractive odor source of two aversive odor sources. The dimension of
+this odor space is 2 (attractive, aversive) despite the number of odor
+sources being 3. The odor sources share a peak intensity of 1. We will
+color the attractive odor source orange and the aversive ones blue.
+
+.. code-block:: ipython3
+   :linenos:
+
+    import numpy as np
+
+    # Odor source: array of shape (num_odor_sources, 3) - xyz coords of odor sources
+    odor_source = np.array([[24, 0, 1.5], [8, -4, 1.5], [16, 4, 1.5]])
+
+    # Peak intensities: array of shape (num_odor_sources, odor_dimesions)
+    # For each odor source, if the intensity is (x, 0) then the odor is in the 1st dimension
+    # (in this case attractive). If it's (0, x) then it's in the 2nd dimension (in this case
+    # aversive)
+    peak_intensity = np.array([[1, 0], [0, 1], [0, 1]])
+
+    # Marker colors: array of shape (num_odor_sources, 4) - RGBA values for each marker,
+    # normalized to [0, 1]
+    marker_colors = [[255, 127, 14], [31, 119, 180], [31, 119, 180]]
+    marker_colors = np.array([[*np.array(color) / 255, 1] for color in marker_colors])
+
+    odor_dimesions = len(peak_intensity[0])
+
+Let’s create the arena using these parameters. The detailed
+documentation of the ``OdorArena`` class can be found in the `API
+reference <https://neuromechfly.org/api_ref/arena.html#flygym.mujoco.arena.OdorArena>`__.
+Its implementation is beyond the scope of this tutorial but can be found
+`here <https://github.com/NeLy-EPFL/flygym/blob/main/flygym/mujoco/arena/sensory_environment.py>`__.
+
+.. code-block:: ipython3
+   :linenos:
+
+    from flygym.mujoco.arena import OdorArena
+
+    arena = OdorArena(
+        odor_source=odor_source,
+        peak_intensity=peak_intensity,
+        diffuse_func=lambda x: x**-2,
+        marker_colors=marker_colors,
+        marker_size=0.3,
+    )
+
+Let’s put a fly in it. As before, we will run a few iterations so it
+stablizes on the ground.
+
+.. code-block:: ipython3
+   :linenos:
+
+    import matplotlib.pyplot as plt
+    from flygym.mujoco import Parameters
+    from flygym.mujoco.examples.turning_controller import HybridTurningNMF
+
+
+    contact_sensor_placements = [
+        f"{leg}{segment}"
+        for leg in ["LF", "LM", "LH", "RF", "RM", "RH"]
+        for segment in ["Tibia", "Tarsus1", "Tarsus2", "Tarsus3", "Tarsus4", "Tarsus5"]
+    ]
+    sim_params = Parameters(
+        timestep=1e-4,
+        render_mode="saved",
+        render_playspeed=0.5,
+        render_window_size=(800, 608),
+        enable_olfaction=True,
+        enable_adhesion=True,
+        draw_adhesion=False,
+        render_camera="birdeye_cam",
+    )
+    sim = HybridTurningNMF(
+        sim_params=sim_params,
+        arena=arena,
+        spawn_pos=(0, 0, 0.2),
+        contact_sensor_placements=contact_sensor_placements,
+    )
+    for i in range(500):
+        sim.step(np.zeros(2))
+        sim.render()
+    fig, ax = plt.subplots(1, 1, figsize=(5, 4), tight_layout=True)
+    ax.imshow(sim._frames[-1])
+    ax.axis("off")
+    fig.savefig("./outputs/olfaction_env.png")
+
+
+
+.. figure :: https://github.com/NeLy-EPFL/_media/blob/main/flygym/olfaction_env.png?raw=true
+   :width: 500
+
+
+
+Controller for odor taxis
+-------------------------
+
+Let’s design a simple hand-tuned controller for odor-guided taxis. We
+start by calculating the left-right asymmetry of the intensity :math:`I`
+for each odor :math:`o`:
+
+.. math::
+
+
+   \Delta I_o = \frac{I_\text{left,o} - I_\text{right,o}}{(I_\text{left,o} + I_\text{right,o}) / 2}
+
+Then, we multiply :math:`\Delta I_o` by a gain :math:`\gamma_o` for each
+odor dimension and take the sum :math:`s`. Attractive and aversive odors
+will have different signs in their gains.
+
+.. math::
+
+
+   s = \sum_{o} \gamma_o \Delta I_o
+
+We transform :math:`s` nonlinearly to avoid turns that are too drastic
+when the asymmetry is subtle and to to crop it to the range [0, 1). This
+gives us a turning bias :math:`b`:
+
+.. math::
+
+
+   b = \tanh(s^2)
+
+Finally, we modulate the descending signal :math:`\delta` based on
+:math:`b` and the sign of :math:`s`:
+
+.. math::
+
+
+   \delta_\text{left} = 
+       \begin{cases}
+       \delta_\text{max} & \text{if } s>0\\
+       \delta_\text{max} - b(\delta_\text{max} - \delta_\text{min})  & \text{otherwise}
+       \end{cases}
+       \qquad
+       \delta_\text{right} = 
+       \begin{cases}
+       \delta_\text{max} - b(\delta_\text{max} - \delta_\text{min}) & \text{if } s>0\\
+       \delta_\text{max}  & \text{otherwise}
+       \end{cases}
+
+where, :math:`\delta_\text{min}`, :math:`\delta_\text{max}` define the
+range of the descending signal. Here, we will use the following
+parameters:
+
+-  :math:`\gamma_\text{attractive} = -500` (negative ipsilateral gain
+   leads to positive taxis)
+-  :math:`\gamma_\text{aversive} = 80` (positive ipsilateral gain leads
+   to negative taxis)
+-  :math:`\delta_\text{min} = 0.2`
+-  :math:`\delta_\text{max} = 1`
+
+As before, we will recalculate the steering signal every 0.05 seconds.
+Let’s implement this in Python:
+
+.. code-block:: ipython3
+   :linenos:
+
+    from tqdm import trange
+
+    attractive_gain = -500
+    aversive_gain = 80
+    decision_interval = 0.05
+    run_time = 5
+    num_decision_steps = int(run_time / decision_interval)
+    physics_steps_per_decision_step = int(decision_interval / sim_params.timestep)
+
+    obs_hist = []
+    odor_history = []
+    obs, _ = sim.reset()
+    for i in trange(num_decision_steps):
+        attractive_intensities = np.average(
+            obs["odor_intensity"][0, :].reshape(2, 2), axis=0, weights=[9, 1]
+        )
+        aversive_intensities = np.average(
+            obs["odor_intensity"][1, :].reshape(2, 2), axis=0, weights=[10, 0]
+        )
+        attractive_bias = (
+            attractive_gain
+            * (attractive_intensities[0] - attractive_intensities[1])
+            / attractive_intensities.mean()
+        )
+        aversive_bias = (
+            aversive_gain
+            * (aversive_intensities[0] - aversive_intensities[1])
+            / aversive_intensities.mean()
+        )
+        effective_bias = aversive_bias + attractive_bias
+        effective_bias_norm = np.tanh(effective_bias**2) * np.sign(effective_bias)
+        assert np.sign(effective_bias_norm) == np.sign(effective_bias)
+
+        control_signal = np.ones((2,))
+        side_to_modulate = int(effective_bias_norm > 0)
+        modulation_amount = np.abs(effective_bias_norm) * 0.8
+        control_signal[side_to_modulate] -= modulation_amount
+
+        for j in range(physics_steps_per_decision_step):
+            obs, _, _, _, _ = sim.step(control_signal)
+            rendered_img = sim.render()
+            if rendered_img is not None:
+                # record odor intensity too for video
+                odor_history.append(obs["odor_intensity"])
+            obs_hist.append(obs)
+
+        # Stop when the fly is within 2mm of the attractive odor source
+        if np.linalg.norm(obs["fly"][0, :2] - odor_source[0, :2]) < 2:
+            break
+
+
+.. parsed-literal::
+
+     77%|███████▋  | 77/100 [01:48<00:32,  1.41s/it]
+
+
+We can visualize the fly trajectory:
+
+.. code-block:: ipython3
+   :linenos:
+
+    fly_pos_hist = np.array([obs["fly"][0, :2] for obs in obs_hist])
+    fig, ax = plt.subplots(1, 1, figsize=(5, 4), tight_layout=True)
+    ax.scatter(
+        [odor_source[0, 0]],
+        [odor_source[0, 1]],
+        marker="o",
+        color="tab:orange",
+        s=50,
+        label="Attractive",
+    )
+    ax.scatter(
+        [odor_source[1, 0]],
+        [odor_source[1, 1]],
+        marker="o",
+        color="tab:blue",
+        s=50,
+        label="Aversive",
+    )
+    ax.scatter([odor_source[2, 0]], [odor_source[2, 1]], marker="o", color="tab:blue", s=50)
+    ax.plot(fly_pos_hist[:, 0], fly_pos_hist[:, 1], color="k", label="Fly trajectory")
+    ax.set_aspect("equal")
+    ax.set_xlim(-1, 25)
+    ax.set_ylim(-5, 5)
+    ax.set_xlabel("x (mm)")
+    ax.set_ylabel("y (mm)")
+    ax.legend(ncols=3, loc="lower center", bbox_to_anchor=(0.5, -0.6))
+    fig.savefig("./outputs/odor_taxis_trajectory.png")
+
+
+
+.. figure :: https://github.com/NeLy-EPFL/_media/blob/main/flygym/odor_taxis_trajectory.png?raw=true
+   :width: 500
+
+
+We can also generate the video:
+
+.. code-block:: ipython3
+   :linenos:
+
+    sim.save_video("./outputs/odor_taxis.mp4")
+
+
+.. raw:: html
+
+   <video src="https://raw.githubusercontent.com/NeLy-EPFL/_media/main/flygym/odor_taxis.mp4" controls="controls" style="max-width: 500px;"></video>