asv_catchment_area.py

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import pkg_resources
from navipy.database import DataBaseLoad
from navipy.processing.pcode import apcv
from navipy.moving.agent import GraphAgent
from navipy import Brain


# 0) Define a class heriting from Brain
class ASVBrain(Brain):
    def __init__(self, renderer=None, channel=0):
        Brain.__init__(self,renderer=renderer)
        # Init memory
        locid = mydb.posorients[(mydb.posorients.x==0) & (mydb.posorients.y==0)].index[0]
        posorient = mydb.posorients.loc[locid, :]
        self.update(posorient)
        self.channel = channel
        self.memory = self.asv()
        

    def asv(self):
        if self.channel >3:
            return apcv(1/self.vision.scene,
                      self.vision.viewing_directions)[..., 3,:]
        else:
            return apcv(self.vision.scene,
                      self.vision.viewing_directions)[..., self.channel,:]
        

    def velocity(self):
        homing_vector = self.memory - self.asv()
        homing_vector = np.squeeze(homing_vector)
        velocity = pd.Series(data=0,
                         index=['dx', 'dy', 'dz',
                                'dalpha_0', 'dalpha_1', 'dalpha_2'])
        velocity[['dx', 'dy', 'dz']] = homing_vector
        return velocity
    
# 1) Connect to the database
mydb_filename = '/home/bolirev/database.db'
mydb = DataBaseLoad(mydb_filename)
for channel  in range(5):
    mybrain = ASVBrain(renderer=mydb, channel=channel)
    # Create a grid agent
    my_agent = GraphAgent(mybrain)
    attractors = my_agent.find_attractors()
    print(attractors)

    f, axarr = plt.subplots(1, 1, figsize=(10, 10))
    axarr.plot(trajectory.x, trajectory.y)
    gridtoplot = mydb.posorients[(mydb.posorients.x>-1) &
                                 (mydb.posorients.y>-1) &
                                 (mydb.posorients.x<1) &
                                 (mydb.posorients.y<1)]
    axarr.plot([0],[0],'ro')
    axarr.plot(gridtoplot.x,gridtoplot.y,'.',color='gray')
    axarr.set_xlim([-10,10])
    axarr.set_ylim([-10,10])
    f.show()
    f.savefig('plots/asv_homing_grid_channel_{}.svg'.format(channel))