Visualizing sensory landscapes

Sensory landscapes of diverse modalities are available :

• Olfactory (Odorscape)

• Temperature (Thermoscape)

• Wind (Windscape)

Initialize the larvaworld registry. This loads some components from disc and builds the rest on the fly.

import numpy as np

from larvaworld.lib import reg, sim, util

# Tutorial safety switches (set True to run interactive demos)
RUN_SINGLE_DEMO = False
RUN_ALL_COMBINATIONS = False
RUN_OLD_METHOD = False
SAVE_MEDIA = False
MEDIA_DIR = "./media"

from larvaworld.lib.param.composition import Odor


def get_odorscape(m):
    if m == "Diffusion":
        oS = reg.gen.DiffusionValueLayer(gaussian_sigma=(0.95, 0.5), evap_const=0.9)
        oR = Odor.oD(id="Odor_R")
        oL = Odor.oD(id="Odor_L")
    elif m == "Gaussian":
        oS = reg.gen.GaussianValueLayer()
        oR = Odor.oG(id="Odor_R")
        oL = Odor.oG(id="Odor_L")
    else:
        raise ValueError("Not implemented")

    kws = {
        "group": "Source",
        "radius": 0.003,
        "amount": 0.0,
    }

    sus = {
        **reg.gen.Food(
            unique_id="Source_L", c="blue", odor=oL, pos=(-0.01, 0.0), **kws
        ).entry(),
        **reg.gen.Food(
            unique_id="Source_R", c="cyan", odor=oR, pos=(0.01, 0.0), **kws
        ).entry(),
    }

    return util.AttrDict(
        {"food_params": reg.gen.FoodConf(source_units=sus), "odorscape": oS}
    )


def get_windscape(m):
    kws = {
        "duration": 5,
        "speed": 50,
    }

    Npuffs = 10
    if m == "single":
        puffs = {
            i: reg.gen.AirPuff(
                direction=i / Npuffs * 2 * np.pi, start_time=5 + 10 * i, **kws
            ).nestedConf
            for i in range(Npuffs)
        }
        ws = 0.0
    elif m == "repetitive":
        puffs = {
            "puff_group": reg.gen.AirPuff(
                direction=np.pi / 4, start_time=5, N=Npuffs, interval=10.0, **kws
            ).nestedConf
        }
        ws = 0.0
    elif m == "no":
        puffs = {}
        ws = 10.0
    else:
        raise ValueError("Not implemented")

    wS = reg.gen.WindScape(wind_direction=0.0, wind_speed=ws, puffs=puffs)

    bl = {"Border": reg.gen.Border(vertices=[(-0.03, 0.02), (0.03, 0.02)])}
    bl = {}
    return util.AttrDict({"windscape": wS, "border_list": bl})


def get_env(Om=None, Pm=None):
    dO = get_odorscape(m=Om) if Om is not None else {}
    dW = get_windscape(m=Pm) if Pm is not None else {}
    return reg.gen.Env(**dO, **dW)


def get_id(Om=None, Pm=None, Wm=None):
    if Om is not None and Pm is None:
        filename = f"{Om}_odorscape"
    elif Om is None and Pm is not None:
        filename = f"{Pm}_air-puffs_variable_wind_{Wm}"
    else:
        filename = f"{Om}_odorscape_{Pm}_air-puffs_variable_wind_{Wm}"
    return filename

# New method


def run_scape(Om=None, Pm=None, Wm=None, duration=0.15, **kwargs):
    def func1(model):
        model.windscape.set_wind_direction((model.t / 10 / np.pi) % (2 * np.pi))

    def func2(model):
        model.windscape.wind_speed = model.t % 100

    dic = {"direction": func1, "speed": func2}

    func = dic[Wm] if Wm in dic.keys() else None

    id = get_id(Om=Om, Pm=Pm, Wm=Wm)
    env = get_env(Om=Om, Pm=Pm)
    env.visualize(
        duration=duration,
        id=id,
        screen_kws={
            "vis_mode": "video",
            "save_video": SAVE_MEDIA,
            "media_dir": MEDIA_DIR,
            "video_file": id,
        },
        func=func,
        **kwargs,
    )

if RUN_SINGLE_DEMO:
    # Run single scape

    run_scape(Om="Diffusion", Pm="repetitive", Wm="speed", duration=2)

if RUN_ALL_COMBINATIONS:
    # Run all scape combinations

    # exp='Wind&Odorscape visualization'
    # exp='dish'

    for Om in ["Gaussian", "Diffusion", None]:
        for Pm in [None, "single", "repetitive", "no"]:
            for Wm in [None, "direction", "speed", "no"]:
                if Pm is None and Wm is not None:
                    continue
                elif Pm is not None and Wm is None:
                    continue
                elif Om is None and Pm is None:
                    continue
                run_scape(Om=Om, Pm=Pm, Wm=Wm)
                # print(conf.screen_kws.video_file)

if RUN_OLD_METHOD:
    # Old method

    def run_scape(Om=None, Pm=None, Wm=None, duration=0.3, **kwargs):
        p = util.AttrDict({"env_params": get_env(Om=Om, Pm=Pm).nestedConf})
        id = get_id(Om=Om, Pm=Pm, Wm=Wm)
        conf = util.AttrDict(
            {
                "id": id,
                "parameters": p,
                "screen_kws": {
                    "vis_mode": "video",
                    "fps": 10,
                    "show_display": True,
                    "save_video": True,
                    "media_dir": "./media",
                    "video_file": id,
                    "odor_aura": True,
                },
            }
        )

        m = sim.base_run.BaseRun(
            **conf, runtype="Exp", experiment="dish", duration=duration, **kwargs
        )
        m.build_env(m.p.env_params)
        # print(m.p.steps)
        m.sim_setup(steps=m.p.steps, seed=None)
        if Om is not None:
            m.odor_layers["Odor_L"].visible = True
            # env.odor_layers['Odor_R'].visible = True
            # env.screen_manager.odor_aura = True
        if Pm is not None:
            m.windscape.visible = True
        while m.running:
            if Wm == "direction":
                m.windscape.set_wind_direction((m.t / 10 / np.pi) % (2 * np.pi))
            elif Wm == "speed":
                m.windscape.wind_speed = m.t % 100
            m.sim_step()
        m.end()
        m.screen_manager.close()

from larvaworld.lib.model.envs import *

(
    GaussianValueLayer.default_color,
    WindScape.default_color,
    DiffusionValueLayer.default_color,
    ThermoScape.default_color,
)

# GaussianValueLayer.__bases__[0].default_color='yellow'

OdorScape().default_color, GaussianValueLayer().default_color
# OdorScape.default_color='grey'