Source code for planktos.swarm

'''
Swarm class of Planktos.

Created on Tues Jan 24 2017

Author: Christopher Strickland

Email: cstric12@utk.edu
'''

import sys, os, warnings
from pathlib import Path
import numpy as np
import numpy.ma as ma
from scipy import interpolate, stats
from scipy.spatial import distance
import pandas as pd
if sys.platform == 'darwin': # OSX backend does not support blitting
    import matplotlib
    matplotlib.use('Qt5Agg')
import matplotlib.pyplot as plt
from matplotlib import animation, colors

from .environment import environment
from . import dataio
from . import motion

__author__ = "Christopher Strickland"
__email__ = "cstric12@utk.edu"
__copyright__ = "Copyright 2017, Christopher Strickland"

class swarm:
    '''
    Fundamental Planktos object describing a group of similar agents.

    The swarm class (alongside the environment class) provides the main agent 
    functionality of Planktos. Each swarm object should be thought of as a group 
    of similar (though not necessarily identical) agents. Planktos implements 
    agents in this way rather than as individual objects for speed purposes; it 
    is easier to vectorize a swarm of agents than individual agent objects, and 
    also much easier to plot them all, get data on them all, etc.

    The swarm object contains all information on the agents' positions, 
    properties, and movement algorithm. It also handles plotting of the agents 
    and saving of agent data to file for further analysis.

    Initial agent velocities will be set as the local fluid velocity if present,
    otherwise zero. Assignment to the velocities attribute can be made directly 
    for other initial conditions.

    NOTE: To customize agent behavior, subclass this class and re-implement the
    method get_positions (do not change the call signature).

    Parameters
    ----------
    swarm_size : int, default=100
        Number of agents in the swarm. ignored when using the 'grid' init method. 
    envir : environment object, optional
        Environment for the swarm to exist in. Defaults to a newly initialized 
        environment with all of the defaults.
    init : {'random', 'grid', ndarray}, default='random'
        Method for initalizing agent positions.
        
        * 'random': Uniform random distribution throughout the domain
        * 'grid': Uniform grid on interior of the domain, including capability
          to leave out closed immersed structures. In this case, swarm_size 
          is ignored since it is determined by the grid dimensions.
          Requires the additional keyword parameters grid_dim and testdir.
        * 1D array: All positions set to a single point given by the x,y,[z] 
          coordinates of this array
        * 2D array: All positions as specified. Shape of array should be NxD, 
          where N is the number of agents and D is spatial dimension. In this 
          case, swarm_size is ignored.
    seed : int, optional
        Seed for random number generator
    shared_props : dictionary, optional
        dictionary of properties shared by all agents as name-value pairs. If 
        none are provided, two default properties will be created, 'mu' and 'cov', 
        corresponding to intrinsic mean drift and a covariance matrix for 
        brownian motion respectively. 'mu' will be set to an array of zeros with 
        length matching the spatial dimension, and 'cov' will be set to an 
        identity matrix of appropriate size according to the spatial dimension. 
        This allows the default agent behavior to be unbiased brownian motion.  
        Examples:  
        * diam: diameter of the particles
        * m: mass of the particles
        * Cd: drag coefficient of the particles
        * cross_sec: cross-sectional area of the particles
        * R: density ratio
    props : Pandas dataframe of individual agent properties, optional
        Pandas dataframe of individual agent properties that vary between agents. 
        This is the method by which individual variation among the agents should 
        be specified. The number of rows in the dataframe should match the 
        number of agents. If no dataframe is supplied, a default one is created 
        which contains only the agent starting positions in a column entitled
        'start_pos'. This is to aid in creating more properties later, if 
        desired, as it is only necessary to add columns to the existing dataframe. 
    **kwargs : dict, optional
        keyword arguments to be used in the 'grid' initialization method or
        values to be set as a swarm object property. In the latter case, these 
        values can be floats, ndarrays, or iterables, but keep in mind that
        problems will result with parsing if the number of agents is
        equal to the spatial dimension - this is to be avoided. This method of 
        specifying agent properties is depreciated: use the shared_props 
        dictionary instead.

    Other Parameters
    ----------------
    grid_dim : tuple of int (x, y, [z])
        number of grid points in x, y, [and z] directions for 'grid' initialization
    testdir : {'x0', 'x1', 'y0', 'y1', ['z0'], ['z1']}, optional
        two character string for testing if grid points are in the interior of 
        an immersed structure and if so, masking them in the grid initialization. 
        The first char is x,y, or z denoting the dimensional direction of the 
        search ray, the second is either 0 or 1 denoting the direction 
        (backward vs. forward) along that direction. See documentation of 
        swarm.grid_init for more information.

    Attributes
    ----------
    envir : environment object
        environment that this swarm belongs to
    positions : masked array, shape Nx2 (2D) or Nx3 (3D)
        spatial location of all the agents in the swarm. the mask is False for 
        any row corresponding to an agent that is within the spatial boundaries 
        of the environment, otherwise the mask for the row is set to True and 
        the position of that agent is no longer updated
    pos_history : list of masked arrays
        all previous position arrays are stored here. to get their corresponding 
        times, check the time_history attribute of the swarm's environment.
    full_pos_history : list of masked arrays
        same as pos_history, but also includes the positions attribute as the 
        last entry in the list
    velocities : masked array, shape Nx2 (2D) or Nx3 (3D)
        velocity of all the agents in the swarm. same masking properties as 
        positions
    accelerations : masked array, shape Nx2 (2D) or Nx3 (3D)
        accelerations of all the agents in the swarm. same masking properties as 
        positions
    ib_collision : 1D array of bool with length equal to the swarm size
        For each agent, True if the agent collided with an immersed boundary in
        the most recent time it moved. False otherwise.
    props : pandas DataFrame
        Pandas dataframe of individual agent properties that vary between agents. 
        This is the method by which individual variation among the agents should 
        be specified.
    shared_props : dictionary
        dictionary of properties shared by all agents as name-value pairs
    rndState : numpy Generator object
        random number generator for this swarm, seeded by the "seed" parameter

    Notes
    -----
    If the agent behavior you are looking for is simply brownian motion with 
    fluid advection, all you need to do is change the 'cov' entry in the 
    shared_props dictionary to a covariance matrix that matches the amount of 
    jitter you are looking for. You can also add fixed directional bias by 
    editing the 'mu' shared_props entry. This default behavior is then 
    accomplished by solving the relevant SDE using Euler steps, where the step 
    size is the dt argument of the move method which you call in a loop, e.g.::
    
        for ii in range(50):
            swm.move(0.1)

    In order to accomodate general, user-defined behavior algorithms, all other 
    agent behaviors should be explicitly specified by subclassing this swarm 
    class and overriding the get_positions method. This is easy, and takes the 
    following form: ::

        class myagents(planktos.swarm):
            
            def get_positions(self, dt, params=None):
                #
                # Put any calculations here that are necessary to determine
                #   where the agents should end up after a time step of length 
                #   dt assuming they don't run into a boundary of any sort. 
                #   Boundary conditions, mesh crossings, etc. will be handled 
                #   automatically by Planktos after this function returns. The
                #   new positions you return should be an ndarray of shape NxD
                #   where N is the number of agents in the swarm and D is the 
                #   spatial dimension of the system. The params argument is 
                #   there in case you want this method to take in any external 
                #   info (e.g. time-varying forcing functions, user-controlled 
                #   behavior switching, etc.). Note that this method has full 
                #   access to all of the swarm attributes via the "self" 
                #   argument. For example, self.positions will return an NxD 
                #   masked array of current agent positions. The one thing this 
                #   method SHOULD NOT do is set the positions, velocities, or 
                #   accelerations attributes of the swarm. This will be handled 
                #   automatically after this method returns, and after boundary 
                #   conditions have been checked.

                return newpositions

    Then, when you create a swarm object, create it using::
    
        swrm = myagents() # add swarm parameters as necessary, as documented above

    This will create a swarm object, but with your my_positions method instead 
    of the default one!

    Examples
    --------
    Create a default swarm in an environment with some fluid data loaded and tiled.

    >>> envir = planktos.environment()
    >>> envir.read_IBAMR3d_vtk_dataset('../tests/IBAMR_test_data', start=5, finish=None)
    >>> envir.tile_flow(3,3)
    >>> swrm = swarm(envir=envir)

    '''

    def __init__(self, swarm_size=100, envir=None, init='random', seed=None, 
                 shared_props=None, props=None, **kwargs):

        # use a new, 3D default environment if one was not given. Or infer
        #   dimension from init if possible.
        if envir is None:
            if isinstance(init,str):
                self.envir = environment(init_swarms=self, Lz=10)
            elif isinstance(init,np.ndarray) and len(init.shape) == 2:
                if init.shape[1] == 2:
                    self.envir = environment(init_swarms=self)
                else:
                    self.envir = environment(init_swarms=self, Lz=10)
            else:
                if len(init) == 2:
                    self.envir = environment(init_swarms=self)
                else:
                    self.envir = environment(init_swarms=self, Lz=10)
        else:
            try:
                assert envir.__class__.__name__ == 'environment'
                envir.swarms.append(self)
                self.envir = envir
            except AssertionError as ae:
                print("Error: invalid environment object.")
                raise ae

        # initialize random number generator
        self.rndState = np.random.default_rng(seed=seed)

        # initialize agent locations
        if isinstance(init,np.ndarray) and len(init.shape) == 2:
            swarm_size = init.shape[0]
        self.positions = ma.zeros((swarm_size, len(self.envir.L)))
        if isinstance(init,str):
            if init == 'random':
                print('Initializing swarm with uniform random positions...')
                for ii in range(len(self.envir.L)):
                    self.positions[:,ii] = self.rndState.uniform(0, 
                                        self.envir.L[ii], self.positions.shape[0])
            elif init == 'grid':
                assert 'grid_dim' in kwargs, "Required key word argument grid_dim missing for grid init."
                x_num = kwargs['grid_dim'][0]; y_num = kwargs['grid_dim'][1]
                if len(self.envir.L) > 2:
                    z_num = kwargs['grid_dim'][2]
                else:
                    z_num = None
                if 'testdir' in kwargs:
                    testdir = kwargs['testdir']
                else:
                    testdir = None
                print('Initializing swarm with grid positions...')
                self.positions = self.grid_init(x_num, y_num, z_num, testdir)
                swarm_size = self.positions.shape[0]
            else:
                print("Initialization method {} not implemented.".format(init))
                print("Exiting...")
                raise NameError
        else:
            if isinstance(init,np.ndarray) and len(init.shape) == 2:
                assert init.shape[1] == len(self.envir.L),\
                    "Initial location data must be Nx{} to match number of agents.".format(
                    len(self.envir.L))
                self.positions[:,:] = init[:,:]
            else:
                for ii in range(len(self.envir.L)):
                    self.positions[:,ii] = init[ii]

        # Due to overloading of the __setattr__ method below, positions, velocities, 
        #   and accelerations should always have a hard mask automatically.

        # initialize agent velocities
        if self.envir.flow is not None:
            self.velocities = ma.array(self.get_fluid_drift(), mask=self.positions.mask.copy())
        else:
            self.velocities = ma.array(np.zeros((swarm_size, len(self.envir.L))), 
                                       mask=self.positions.mask.copy())

        # initialize agent accelerations
        self.accelerations = ma.array(np.zeros((swarm_size, len(self.envir.L))),
                                      mask=self.positions.mask.copy())

        # Initialize position history
        self.pos_history = []

        # Apply boundary conditions in case of domain mismatch
        self.apply_boundary_conditions(ib_collisions=None)

        # Initialize IB collision detection
        self.ib_collision = np.full(swarm_size, False)

        # initialize Dataframe of non-shared properties
        if props is None:
            self.props = pd.DataFrame()
            # with random cov
            # self.props = pd.DataFrame(
            #     {'start_pos': [tuple(self.positions[ii,:]) for ii in range(swarm_size)],
            #     'cov': [np.eye(len(self.envir.L))*(0.5+np.random.rand()) for ii in range(swarm_size)]}
            # )
        else:
            self.props = props

        # Dictionary of shared properties
        if shared_props is None:
            self.shared_props = {}
        else:
            self.shared_props = shared_props

        # Include parameters for a uniform standard random walk by default
        if 'mu' not in self.shared_props and 'mu' not in self.props:
            self.shared_props['mu'] = np.zeros(len(self.envir.L))
        if 'cov' not in self.shared_props and 'cov' not in self.props:
            self.shared_props['cov'] = np.eye(len(self.envir.L))

        # Record any kwargs as swarm parameters
        for name, obj in kwargs.items():
            try:
                if isinstance(obj,np.ndarray) and obj.shape[0] == swarm_size and\
                    obj.shape[0] != len(self.envir.L):
                    self.props[name] = obj
                elif isinstance(obj,np.ndarray):
                    self.shared_props[name] = obj
                elif len(obj) == swarm_size and len(obj) != len(self.envir.L):
                    self.props[name] = obj
                else:
                    # convert iterable to ndarray first
                    self.shared_props[name] = np.array(obj)
            except TypeError:
                # Called len on something that wasn't iterable
                self.shared_props[name] = obj
                    


    # Make sure mask is always hardened for positions, velocities, and accelerations
    def __setattr__(self, name, value):
        if name in ['positions', 'velocities', 'accelerations']:
            assert isinstance(value, np.ndarray), name+" must be an array or masked array."
            if not isinstance(value, ma.masked_array):
                value = ma.masked_array(value)
            value.harden_mask()
        super().__setattr__(name, value)



[docs] def grid_init(self, x_num, y_num, z_num=None, testdir=None): '''Return a flattened array which describes a regular grid of locations, except potentially masking any grid points in the interior of a closed, immersed structure. The full, unmasked grid will be x_num by y_num [by z_num] on the interior and boundaries of the domain. The output of this method is appropriate for finding FTLE, and that is its main purpose. It will automatically be called by the environment class's calculate_FTLE method, and if you want to initialize a swarm with a grid this is possible by passing the init='grid' keyword argument when the swarm is created. So there is probably no reason to use this method directly. Grid list moves in the [Z direction], Y direction, then X direction (due to C order of memory layout). Parameters ---------- x_num, y_num, [z_num] : int number of grid points in each direction testdir : {'x0', 'x1', 'y0', 'y1', ['z0'], ['z1']}, optional to check if a point is an interior to an immersed structure, a line will be drawn from the point to a domain boundary. If the number of immersed boundary intersections is odd, the point will be considered interior and masked. This check will not be run at all if testdir is None. Otherwise, specify a direction with one of the following: 'x0','x1','y0','y1','z0','z1' (the last two for 3D problems only) denoting the dimension (x,y, or z) and the direction (0 for negative, 1 for positive). Notes ----- This algorithm is meant as a huristic only! It is not guaranteed to mask all interior grid points, and will mask non-interior points if there is not a clear line from the point to one of the boundaries of the domain. If this method fails for your geometry and better accuracy is needed, use this method as a starting point and mask/unmask as necessary. ''' # Form initial grid x_pts = np.linspace(0, self.envir.L[0], x_num) y_pts = np.linspace(0, self.envir.L[1], y_num) if z_num is not None: z_pts = np.linspace(0, self.envir.L[2], z_num) X1, X2, X3 = np.meshgrid(x_pts, y_pts, z_pts, indexing='ij') xidx, yidx, zidx = np.meshgrid(np.arange(x_num), np.arange(y_num), np.arange(z_num), indexing='ij') DIM = 3 elif len(self.envir.L) > 2: raise RuntimeError("Must specify z_num for 3D problems.") else: X1, X2 = np.meshgrid(x_pts, y_pts, indexing='ij') xidx, yidx = np.meshgrid(np.arange(x_num), np.arange(y_num), indexing='ij') DIM = 2 if testdir is None: if DIM == 2: return ma.array([X1.flatten(), X2.flatten()]).T else: return ma.array([X1.flatten(), X2.flatten(), X3.flatten]).T elif testdir[0] == 'z' and len(self.envir.L) < 3: raise RuntimeError("z-direction unavailable in 2D problems.") # Translate directional input startdim = [0,0] if testdir[0] == 'x': startdim[0] = 0 if DIM == 2: perp_idx = 1 else: perp_idx = [1,2] elif testdir[0] == 'y': startdim[0] = 1 if DIM == 2: perp_idx = 0 else: perp_idx = [0,2] elif testdir[0] == 'z': startdim[0] = 2 perp_idx = [0,1] else: raise RuntimeError("Unrecognized value in testdir, {}.".format(testdir)) try: startdim[1] = int(testdir[1]) - 1 except ValueError: raise RuntimeError("Unrecognized value in testdir, {}.".format(testdir)) # Idea: start on opposite side of domain as given direction and take a # full grid on the boundary. See if there are intersections. # If none, none of the points along that ray need to be tested further # Otherwise, we are also given the intersection points. Use to # deduce the rest. # startdim gives the dimension and index on which to place a grid # Convert X1, X2, X3 to masked arrays X1 = ma.array(X1) X2 = ma.array(X2) if DIM == 3: X3 = ma.array(X3) grids = [X1, X2, X3] idx_list = [xidx, yidx, zidx] else: grids = [X1, X2] idx_list = [xidx, yidx] # get a list of the gridpoints on correct side of the domain firstpts = [] first_idx_list = [] for X, idx in zip(grids,idx_list): if startdim[0] == 0: firstpts.append(X[startdim[1],...]) first_idx_list.append(idx[startdim[1],...]) elif startdim[0] == 1: firstpts.append(X[:,startdim[1],...]) first_idx_list.append(idx[:,startdim[1],...]) else: firstpts.append(X[:,:,startdim[1]]) first_idx_list.append(idx[:,:,startdim[1]]) firstpts = np.array([X.flatten() for X in firstpts]).T idx_vals = np.array([idx.flatten() for idx in first_idx_list]).T mesh = self.envir.ibmesh meshptlist = mesh.reshape((mesh.shape[0]*mesh.shape[1],mesh.shape[2])) for pt, idx in zip(firstpts,idx_vals): # for each pt in the grid, get a list of eligibible mesh elements as # those who have a point within a cylinder of diameter envir.max_meshpt_dist if DIM == 3: pt_bool = np.linalg.norm(meshptlist[:,perp_idx]-pt[perp_idx], axis=1)<=self.envir.max_meshpt_dist/2 else: pt_bool = np.abs(meshptlist[:,perp_idx]-pt[perp_idx])\ <=self.envir.max_meshpt_dist/2 pt_bool = pt_bool.reshape((mesh.shape[0], mesh.shape[1])) close_mesh = mesh[np.any(pt_bool, axis=1)] endpt = np.array(pt) if startdim[1] == -1: endpt[startdim[0]] = 0 else: endpt[startdim[0]] = self.envir.L[startdim[0]] # Get intersections if DIM == 2: intersections = swarm._seg_intersect_2D(pt, endpt, close_mesh[:,0,:], close_mesh[:,1,:], get_all=True) else: intersections = swarm._seg_intersect_3D_triangles(pt, endpt, close_mesh[:,0,:], close_mesh[:,1,:], close_mesh[:,2,:], get_all=True) # For completeness, we should also worry about edge cases where # intersections are not of mesh facets but of triangle points, but # as a heuristic, we will ignore this. A tweaking of the number of # grid points used could fix this problem in most cases, or it # could be fixed by hand. if intersections is not None: # Sort the intersections by distance away from pt intersections = sorted(intersections, key=lambda x: x[1]) # get list of all x,y, or z values for points along the ray # (where the dimension matches the direction of the ray) if startdim[0] == 0: current_pt_val = pt[0] - 10e-7 if DIM == 3: val_list = X1[:,idx[1],idx[2]] else: val_list = X1[:,idx[1]] elif startdim[0] == 1: current_pt_val = pt[1] - 10e-7 if DIM == 3: val_list = X2[idx[0],:,idx[2]] else: val_list = X2[idx[0],:] else: current_pt_val = pt[2] - 10e-7 val_list = X3[idx[0],idx[1],:] while len(intersections) > 0: n = len(intersections) intersection = intersections.pop(0) if startdim[0] == 0: intersect_val = intersection[0][0] elif startdim[0] == 1: intersect_val = intersection[0][1] else: intersect_val = intersection[0][2] if current_pt_val < intersect_val: pair = [current_pt_val, intersect_val] else: pair = [intersect_val, current_pt_val] # gather all points between current one and intersection # This will always mask points exactly on a mesh boundary bool_list = np.logical_and(pair[0]<=val_list,val_list<=pair[1]) # set mask if number of intersections is odd if n%2 == 1: if startdim[0] == 0: if DIM == 3: X1[bool_list,idx[1],idx[2]] = ma.masked X2[bool_list,idx[1],idx[2]] = ma.masked X3[bool_list,idx[1],idx[2]] = ma.masked else: X1[bool_list,idx[1]] = ma.masked X2[bool_list,idx[1]] = ma.masked elif startdim[0] == 1: if DIM == 3: X1[idx[0],bool_list,idx[2]] = ma.masked X2[idx[0],bool_list,idx[2]] = ma.masked X3[idx[0],bool_list,idx[2]] = ma.masked else: X1[idx[0],bool_list] = ma.masked X2[idx[0],bool_list] = ma.masked else: X1[idx[0],idx[1],bool_list] = ma.masked X2[idx[0],idx[1],bool_list] = ma.masked X3[idx[0],idx[1],bool_list] = ma.masked # Update current_pt_val to latest intersection current_pt_val = intersect_val # all points done. if DIM == 2: return ma.array([X1.flatten(), X2.flatten()]).T else: return ma.array([X1.flatten(), X2.flatten(), X3.flatten()]).T
@property def full_pos_history(self): '''History of self.positions, including present time.''' return [*self.pos_history, self.positions]
[docs] def save_data(self, path, name, pos_fmt='%.18e'): '''Save the full position history (with mask and time stamps) along with current velocity and acceleration to csv files. Save shared_props to a npz file and save props to json. The output format for the position csv is the same as for the save_pos_to_csv method. shared_props is saved as an npz file since it is likely to contain some mixture of scalars and arrays, but does not vary between the agents so is less likely to be loaded outside of Python. props is saved to json since it is likely to contain a variety of types of data, may need to be loaded outside of Python, and json will be human readable. Parameters ---------- path : str directory for storing data name : str prefix name for data files pos_fmt : str format, default='%.18e' format and precision for storing position, vel, and accel data See Also -------- save_pos_to_csv save_pos_to_vtk ''' path = Path(path) if not path.is_dir(): os.makedirs(path) self.save_pos_to_csv(str(path/name), pos_fmt, sv_vel=True, sv_accel=True) props_file = path/(name+'_props.json') self.props.to_json(str(props_file)) shared_props_file = path/(name+'_shared_props.npz') np.savez(str(shared_props_file), **self.shared_props)
[docs] def save_pos_to_csv(self, filename, fmt='%.18e', sv_vel=False, sv_accel=False): '''Save the full position history including present time, with mask and time stamps, to a csv. The output format for the position csv will be as follows: * The first row contains cycle and time information. The cycle is given, and then each time stamp is repeated D times, where D is the spatial dimension of the system. * Each subsequent row corresponds to a different agent in the swarm. * Reading across the columns of an agent row: first, a boolean is given showing the state of the mask for that time step. Agents are masked when they have exited the domain. Then, the position vector is given as a group of D columns for the x, y, (and z) direction. Each set of 1+D columns then corresponds to a different cycle/time, as labeled by the values in the first row. The result is a csv that is N+1 by (1+D)*T, where N is the number of agents, D is the dimension of the system, and T is the number of times recorded. Parameters ---------- filename : str path/name of the file to save the data to fmt : str format, default='%.18e' fmt argument to be passed to numpy.savetxt for format and precision of numerical data sv_vel : bool, default=False whether or not to save the current time velocity data sv_accel : book, default=False whether or not to save the current time acceleration data See Also -------- save_data save_pos_to_vtk ''' if filename[-4:] != '.csv': filename = filename + '.csv' full_time = [*self.envir.time_history, self.envir.time] time_row = np.concatenate([[ii]+[jj]*self.positions.shape[1] for ii,jj in zip(range(len(full_time)), full_time)]) fmtlist = ['%u'] + [fmt]*self.positions.shape[1] np.savetxt(filename, np.vstack((time_row, np.column_stack([mat for pos in self.full_pos_history for mat in (ma.getmaskarray(pos[:,0]), pos.data)]))), fmt=fmtlist*len(full_time), delimiter=',') if sv_vel: vel_filename = filename[:-4] + '_vel.csv' np.savetxt(vel_filename, np.column_stack((self.velocities[:,0].mask, self.velocities.data)), fmt=fmtlist, delimiter=',') if sv_accel: accel_filename = filename[:-4] + '_accel.csv' np.savetxt(accel_filename, np.column_stack((self.accelerations[:,0].mask, self.accelerations.data)), fmt=fmtlist, delimiter=',')
[docs] def save_pos_to_vtk(self, path, name, all=True): '''Save position data to vtk as point data (PolyData). A different file will be created for each time step in the history, or just one file of the current positions will be created if the all argument is False. Parameters ---------- path : str location to save the data name : str name of dataset all : bool if True, save the entire history including the current time. If false, save only the current time. See Also -------- save_data save_pos_to_csv ''' if len(self.envir.L) == 2: DIM2 = True else: DIM2 = False if not all or len(self.envir.time_history) == 0: if DIM2: data = np.zeros((self.positions[~ma.getmaskarray(self.positions[:,0]),:].shape[0],3)) data[:,:2] = self.positions[~ma.getmaskarray(self.positions[:,0]),:] dataio.write_vtk_point_data(path, name, data) else: dataio.write_vtk_point_data(path, name, self.positions[~ma.getmaskarray(self.positions[:,0]),:]) else: for cyc, time in enumerate(self.envir.time_history): if DIM2: data = np.zeros((self.pos_history[cyc][~self.pos_history[cyc][:,0].mask,:].shape[0],3)) data[:,:2] = self.pos_history[cyc][~self.pos_history[cyc][:,0].mask,:].squeeze() dataio.write_vtk_point_data(path, name, data, cycle=cyc, time=time) else: dataio.write_vtk_point_data(path, name, self.pos_history[cyc][~self.pos_history[cyc][:,0].mask,:].squeeze(), cycle=cyc, time=time) cyc = len(self.envir.time_history) if DIM2: data = np.zeros((self.positions[~ma.getmaskarray(self.positions[:,0]),:].shape[0],3)) data[:,:2] = self.positions[~ma.getmaskarray(self.positions[:,0]),:] dataio.write_vtk_point_data(path, name, data, cycle=cyc, time=self.envir.time) else: dataio.write_vtk_point_data(path, name, self.positions[~ma.getmaskarray(self.positions[:,0]),:], cycle=cyc, time=self.envir.time)
def _change_envir(self, envir): '''Manages a change from one environment to another''' if self.positions.shape[1] != len(envir.L): if self.positions.shape[1] > len(envir.L): # Project swarm down to 2D self.positions = self.positions[:,:2] self.velocities = self.velocities[:,:2] self.accelerations = self.accelerations[:,:2] # Update known properties if 'mu' in self.shared_props: self.shared_props['mu'] = self.shared_props['mu'][:2] print('mu has been projected to 2D.') if 'mu' in self.props: for n,mu in enumerate(self.props['mu']): self.props['mu'][n] = mu[:2] print('mu has been projected to 2D.') if 'cov' in self.shared_props: self.shared_props['cov'] = self.shared_props['cov'][:2,:2] print('cov has been projected to 2D.') if 'cov' in self.props: for n,cov in enumerate(self.props['cov']): self.props['cov'][n] = cov[:2,:2] print('cov has been projected to 2D.') # warn about others other_props = [x for x in self.props if x not in ['mu', 'cov']] other_props += [x for x in self.shared_props if x not in ['mu', 'cov']] if len(other_props) > 0: print('WARNING: other properties {} were not projected.'.format(other_props)) else: raise RuntimeError("Swarm dimension smaller than new environment dimension!"+ " Cannot scale up!") self.envir = envir envir.swarms.append(self)
[docs] def calc_re(self, u, diam=None): '''Calculate and return the Reynolds number as experienced by a swarm with characteristic length 'diam' in a fluid moving with velocity u. All other parameters will be pulled from the environment's attributes. If diam is not specified, this method will look for it in the shared_props dictionary of this swarm. Parameters ---------- u : float characteristic fluid speed, m/s diam : float, optional characteristic length scale of a single agent, m Returns ------- float Reynolds number ''' if diam is None: diam = self.shared_props['diam'] else: diam = diam if self.envir.rho is not None and self.envir.mu is not None and\ 'diam' in self.shared_props: return self.envir.rho*u*diam/self.envir.mu else: raise RuntimeError("Parameters necessary for Re calculation in environment are undefined.")
[docs] def move(self, dt=1.0, params=None, ib_collisions='sliding', update_time=True, silent=False): '''Move all organisms in the swarm over one time step of length dt. DO NOT override this method when subclassing; override get_positions instead!!! Performs a lot of utility tasks such as updating the positions and pos_history attributes, checking boundary conditions, and recalculating the current velocities and accelerations attributes. Parameters ---------- dt : float length of time step to move all agents params : any, optional parameters to pass along to get_positions, if necessary ib_collisions : {None, 'sliding' (default), 'sticky'} Type of interaction with immersed boundaries. If None, turn off all interaction with immersed boundaries. In sliding collisions, conduct recursive vector projection until the length of the original vector is exhausted. In sticky collisions, just return the point of intersection. update_time : bool, default=True whether or not to update the environment's time by dt. Probably The only reason to change this to False is if there are multiple swarm objects in the same environment - then you want to update each before incrementing the time in the environment. silent : bool, default=False If True, suppress printing the updated time. See Also -------- get_positions : method that returns (but does not assign) the new positions of the swarm after the time step dt, which Planktos users override in order to specify their own, custom agent behavior. ''' # Put current position in the history self.pos_history.append(self.positions.copy()) # Check that something is left in the domain to move, and move it. if not np.all(self.positions.mask): # Update positions, preserving mask self.positions[:,:] = self.get_positions(dt, params) # Update velocity and acceleration of swarm velocity = (self.positions - self.pos_history[-1])/dt self.accelerations[:,:] = (velocity - self.velocities)/dt self.velocities[:,:] = velocity # Apply boundary conditions. self.apply_boundary_conditions(ib_collisions=ib_collisions) # Record new time if update_time: self.envir.time_history.append(self.envir.time) self.envir.time += dt if not silent: print('time = {}'.format(np.round(self.envir.time,11))) # Check for other swarms in environment and freeze them warned = False for s in self.envir.swarms: if s is not self and len(s.pos_history) < len(self.pos_history): s.pos_history.append(s.positions.copy()) if not warned: warnings.warn("Other swarms in the environment were not"+ " moved during this environmental timestep.\n"+ "If this was not your intent, call"+ " envir.move_swarms instead of this method"+ " to move all the swarms at once.", UserWarning) warned = True
[docs] def get_positions(self, dt, params=None): '''Returns the new agent positions after a time step of dt. THIS IS THE METHOD TO OVERRIDE IF YOU WANT DIFFERENT MOVEMENT! Do not change the call signature. This method returns the new positions of all agents following a time step of length dt, whether due to behavior, drift, or anything else. It should not set the self.positions attribute. Similarly, self.velocities and self.accelerations will automatically be updated outside of this method using finite differences. The only attributes it should change is if there are any user-defined, time-varying agent properties that should be different after the time step (whether shared among all agents, and thus in self.shared_props, or individual to each agent, and thus in self.props). These can be altered directly or by using the add_prop method of this class. In this default implementation, movement is a random walk with drift as given by an Euler step solver of the appropriate SDE for this process. Drift is the local fluid velocity plus self.get_prop('mu') ('mu' is a shared_prop attribute), and the stochasticity is determined by the covariance matrix self.get_prop('cov') ('cov' is also a shared_prop attribute). Parameters ---------- dt : float length of time step params : any, optional any other parameters necessary Returns ------- ndarray : NxD array of new agent positions after a time step of dt given that the agents started at self.positions. N is the number of agents and D is the spatial dimension of the system. Notes ----- When writing code for this method, it can be helpful to make use of the ode generators and solvers in the planktos.motion module. Please see the documentation for the functions of this module for options. To access the current positions of each agent, use self.positions. self.positions is a masked, NxD array of agent positions where the mask refers to whether or not the agent has exited the domain. You do not want to accidently edit self.positions directly, so make sure that you get a value copy of self.positions using self.positions.copy() whenever that copy will be modified. Direct assignment of self.positions is by reference. Similarly,self.velocities and self.accelerations will provide initial velocities and accelerations for the time step for each agent respectively. Use .copy() as necessary and do not directly assign to these variables; they will be automatically updated later in the movement process. The get_fluid_drift method will return the fluid velocity at each agent location using interpolation. Call it once outside of a loop for speed. Similarly, the get_dudt method will return the time derivative of the fluid velocity at the location of each agent. The get_fluid_mag_gradient method will return the gradient of the magnitude of the fluid velocity at the location of each agent. See Also -------- get_prop : given an agent/swarm property name, return the value(s). When accessing a property in swarm.props, this can be preferred over accessing the property directly through the because instead of returning a pandas Series object (for a column in the DataFrame), it automatically converts to a numpy array first. add_prop : add a new agent/swarm property or overwrite an old one get_fluid_drift : return the fluid velocity at each agent location get_dudt : return time derivative of fluid velocity at each agent get_fluid_mag_gradient : return the gradient of the magnitude of the fluid velocity at each agent ''' # default behavior for Euler_brownian_motion is dift due to mu property # plus local fluid velocity and diffusion given by cov property # specifying the covariance matrix. return motion.Euler_brownian_motion(self, dt)
[docs] def get_prop(self, prop_name): '''Return the property requested as either a scalar (if shared) or a numpy array, ready for use in vectorized operations (left-most index specifies the agent). Parameters ---------- prop_name : str name of the property to return Returns ------- property : float or ndarray ''' if prop_name in self.props: if prop_name in self.shared_props: warnings.warn('Property {} exists '.format(prop_name)+ 'in both props and shared_props. Using the props version.') return np.stack(self.props[prop_name].array, axis=0).squeeze() elif prop_name in self.shared_props: return self.shared_props[prop_name] else: raise KeyError('Property {} not found.'.format(prop_name))
[docs] def add_prop(self, prop_name, value, shared=False): '''Method that will automatically delete any conflicting properties when adding a new one. Parameters ---------- prop_name : str name of the property to add value : any value to set the property at shared : bool if False, set as a property that applies to all agents in the swarm. if True, value should be an ndarray with a number of rows equal to the number of agents in the swarm, and the property will be set as a column in the swarm.props DataFrame. ''' if shared: self.shared_props[prop_name] = value if prop_name in self.props: del self.props[prop_name] else: self.props[prop_name] = value if prop_name in self.shared_props: del self.shared_props[prop_name]
[docs] def get_fluid_drift(self, time=None, positions=None): '''Return fluid-based drift for all agents via interpolation. Current swarm position is used unless alternative positions are explicitly passed in. Any passed-in positions must be an NxD array where N is the number of points and D is the spatial dimension of the system. In the returned 2D ndarray, each row corresponds to an agent (in the same order as listed in self.positions) and each column is a dimension. Parameters ---------- time : float, optional time at which to return the fluid drift. defaults to the current environment time positions : ndarray, optional positions at which to return the fluid drift. defaults to the locations of the swarm agents, self.positions Returns ------- ndarray with shape NxD, where N is the number of agents and D the spatial dimension ''' # 3D? DIM3 = (len(self.envir.L) == 3) if positions is None: positions = self.positions # Interpolate fluid flow if self.envir.flow is None: return np.zeros(self.positions.shape) else: if time is None: return self.envir.interpolate_flow(positions, method='linear') else: return self.envir.interpolate_flow(positions, time=time, method='linear')
[docs] def get_dudt(self, time=None, positions=None): '''Return fluid time derivative at given positions via interpolation. Current swarm position is used unless alternative positions are explicitly passed in. In the returned 2D ndarray, each row corresponds to an agent (in the same order as listed in self.positions) and each column is a dimension. Parameters ---------- time : float, optional time at which to return the data. defaults to the current environment time positions : ndarray, optional positions at which to return the data. defaults to the locations of the swarm agents, self.positions Returns ------- ndarray with shape NxD, where N is the number of agents and D the spatial dimension ''' if positions is None: positions = self.positions return self.envir.interpolate_flow(positions, self.envir.dudt(time=time), method='linear')
[docs] def get_fluid_mag_gradient(self, positions=None): '''Return the gradient of the magnitude of the fluid velocity at all agent positions (or at provided positions) via linear interpolation of the gradient. The gradient is linearly interpolated from the fluid grid to the agent locations. The current environment time is always used, interpolated from data if necessary Parameters ---------- positions : ndarray, optional positions at which to return the data. defaults to the locations of the swarm agents, self.positions Returns ------- ndarray with shape NxD, where N is the number of agents and D the spatial dimension ''' if positions is None: positions = self.positions TIME_DEP = len(self.envir.flow[0].shape) != len(self.envir.L) flow_grad = None # If available, use the already calculuated gradient (if it's at the # correct time) if self.envir.mag_grad is not None: if not TIME_DEP: flow_grad = self.envir.mag_grad elif self.envir.mag_grad_time == self.envir.time: flow_grad = self.envir.mag_grad # Otherwise, calculate the gradient if flow_grad is None: self.envir.calculate_mag_gradient() flow_grad = self.envir.mag_grad # Interpolate the gradient at agent positions and return return self.envir.interpolate_flow(positions, flow_grad, method='linear')
[docs] def get_DuDt(self, time=None, positions=None): '''Return the material derivative with respect to time of the fluid velocity at all agent positions (or at provided positions) via linear interpolation of the material gradient. Current swarm position is used unless alternative positions are explicitly passed in. In the returned 2D ndarray, each row corresponds to an agent (in the same order as listed in self.positions) and each column is a dimension. Parameters ---------- time : float, optional time at which to return the data. defaults to the current environment time positions : ndarray, optional positions at which to return the data. defaults to the locations of the swarm agents, self.positions Returns ------- ndarray with shape NxD, where N is the number of agents and D the spatial dimension ''' if positions is None: positions = self.positions if time is None: time = self.envir.time # Calculate if necessary, otherwise use cached copy if self.envir.DuDt is None or self.envir.DuDt_time != time: self.envir.calculate_DuDt(time=time) # Interpolate at agent positions and return return self.envir.interpolate_flow(positions, self.envir.DuDt, method='linear')
[docs] def apply_boundary_conditions(self, ib_collisions='sliding'): '''Apply boundary conditions to self.positions. There should be no reason to call this method directly; it is automatically called by self.move after updating agent positions according to the algorithm found in self.get_positions. This method compares current agent positions (self.positions) to the previous agent positions (last entry in self.pos_history) in order to first: determine if the agent collided with any immersed structures and if so, to update self.positions using a sliding collision algorithm based on vector projection and second: assess whether or not any agents exited the domain and if so, update their positions based on the boundary conditions as specified in the enviornment class (self.envir). For noflux boundary conditions such sliding projections are really simple (since the domain is just a box), so we just do them directly/manually instead of folding them into the far more complex, recursive algorithm used for internal mesh structures. Periodic boundary conditions will recursively check for immersed boundary crossings after each crossing of the domain boundary. Parameters ---------- ib_collisions : {None, 'sliding' (default), 'sticky'} Type of interaction with immersed boundaries. If None, turn off all interaction with immersed boundaries. In sliding collisions, conduct recursive vector projection until the length of the original vector is exhausted. In sticky collisions, just return the point of intersection. ''' ##### Immersed mesh boundaries go first ##### if self.envir.ibmesh is not None and ib_collisions is not None: # Routine for checking IB def IBC_routine(n, self, startpt, endpt, ib_collisions): new_loc = self._apply_internal_BC(startpt, endpt, self.envir.ibmesh, self.envir.max_meshpt_dist, ib_collisions=ib_collisions) self.positions[n] = new_loc if np.any(new_loc != endpt): self.ib_collision[n] = True else: self.ib_collision[n] = False # if there are any masked agents, skip them in the loop if np.any(self.positions.mask): for n, startpt, endpt in \ zip(np.arange(self.positions.shape[0])[~ma.getmaskarray(self.positions[:,0])], self.pos_history[-1][~ma.getmaskarray(self.positions[:,0]),:].copy(), self.positions[~ma.getmaskarray(self.positions[:,0]),:].copy() ): IBC_routine(n, self, startpt, endpt, ib_collisions) # if all are masked, skip all boundary checks elif np.all(self.positions.mask): return # no masked agents: go through all of them else: for n in range(self.positions.shape[0]): startpt = self.pos_history[-1][n,:].copy() endpt = self.positions[n,:].copy() IBC_routine(n, self, startpt, endpt, ib_collisions) ##### Environment Boundary Conditions ##### self._domain_BC_loop(ib_collisions=ib_collisions)
def _domain_BC_loop(self, ib_collisions, idx_array=None): '''Loop over domain boundaries enforcing boundary conditions. Only agents in idx_array will be checked, or all unmasked agents if idx_array is not given. ''' if idx_array is None: idx_array = np.arange(self.positions.shape[0]) status_BC = np.zeros((len(idx_array),len(self.envir.L))) ##### Mark all domain exits! -1 for left, 1 for right ##### # skip masked entries if np.all(self.positions[idx_array].mask): return for dim in range(len(self.envir.L)): leftrow = np.logical_and(self.positions[idx_array,dim] < 0, ~self.positions[idx_array,dim].mask) rightrow = np.logical_and(self.positions[idx_array,dim] > self.envir.L[dim], ~self.positions[idx_array,dim].mask) status_BC[leftrow,dim] = -1 status_BC[rightrow,dim] = 1 ##### In cases where there are multiple exits, find the first ##### BC_mult_bool = np.sum(np.abs(status_BC), axis=1) > 1 if np.any(BC_mult_bool): # mark these for recursion mult_idx = idx_array[BC_mult_bool] # figure out which exit crossing occured first and treat that as the # only one. Use the velocity to parameterize the movement. s_array = np.zeros((len(BC_mult_bool),len(self.envir.L))) for dim in range(len(self.envir.L)): # get multiple crossing entries that have crossing in this dim dim_bool = np.logical_and(BC_mult_bool, status_BC[:,dim] != 0) #full length dim_idx = idx_array[dim_bool] #reduced length right_1 = 1*(status_BC[dim_bool,dim] == 1) #reduced length s_array[dim_bool,dim] = (self.positions[dim_idx,dim]-right_1* self.envir.L[dim])/self.velocities[dim_idx,dim] # for each row in s_array, the dimension with the largest value is # now the one crossed first. first_dim = np.argmax(s_array[BC_mult_bool,:], axis=1) # remove the other crossings from the status_BC array status_vals = status_BC[BC_mult_bool, first_dim] status_BC[BC_mult_bool,:] = 0 status_BC[BC_mult_bool, first_dim] = status_vals BC_bool = np.sum(np.abs(status_BC), axis=1) != 0 BC_bool_check = np.sum(np.abs(status_BC), axis=1) == 1 assert np.all(BC_bool == BC_bool_check), "Some multi crossings left over...?" else: mult_idx = None BC_bool = np.sum(np.abs(status_BC), axis=1) == 1 if not np.any(BC_bool): return ##### Routine for checking IB in periodic case ##### def IBC_routine(idx, self, startpt, endpt, ib_collisions): new_loc = self._apply_internal_BC(startpt, endpt, self.envir.ibmesh, self.envir.max_meshpt_dist, ib_collisions=ib_collisions) self.positions[idx] = new_loc if np.any(new_loc != endpt): self.ib_collision[idx] = True else: self.ib_collision[idx] = False ##### Now apply BC to the first/only boundary crossing ##### for dim, bndry in enumerate(self.envir.bndry): # Check for 3D if dim == 2 and len(self.envir.L) == 2: # Ignore last bndry condition; 2D environment. break ### Left boundary ### left_bool = np.logical_and(BC_bool, status_BC[:,dim]<0) left_idx = idx_array[left_bool] if len(left_idx) > 0: if bndry[0] == 'zero': # mask anything that exited on the left. self.positions[left_idx, :] = ma.masked self.velocities[left_idx, :] = ma.masked self.accelerations[left_idx, :] = ma.masked # no further BC checks are made: masked entries are skipped elif bndry[0] == 'noflux': # agent slides along flat boundary. pos/vel/accel in dir of # boundary will be zero. # additional IB crossings are possible, so first find # point of intersection with the boundary to enable this # check. if self.envir.ibmesh is not None and ib_collisions is not None: s_array = (self.positions[left_idx,dim]-0)/ \ self.velocities[left_idx,dim] startpts = self.positions[left_idx,:] - (np.tile(s_array, (self.velocities.shape[1],1)).T* self.velocities[left_idx,:]) # now update pos/vel/accel self.positions[left_idx, dim] = 0 self.velocities[left_idx, dim] = 0 self.accelerations[left_idx, dim] = 0 # now check for IB crossings. However, due to potential # complex interactions with the noflux boundary, enforce # sticky ib collisions in all cases. if self.envir.ibmesh is not None and ib_collisions is not None: for n, idx in enumerate(left_idx): startpt = startpts[n] endpt = self.positions[idx,:].copy() IBC_routine(idx, self, startpt, endpt, 'sticky') # further domain crossings remain possible elif bndry[0] == 'periodic': # wrap everything exiting on the left to the right self.positions[left_idx, dim] =\ np.mod(self.positions[left_idx, dim],self.envir.L[dim]) # check for IB crossings. first, get the point of re-entry # using the velocity. if self.envir.ibmesh is not None and ib_collisions is not None: s_array = (self.positions[left_idx,dim]- self.envir.L[dim])/self.velocities[left_idx,dim] for n, idx in enumerate(left_idx): startpt = self.positions[idx,:] - \ s_array[n]*self.velocities[idx,:] endpt = self.positions[idx,:].copy() IBC_routine(idx, self, startpt, endpt, ib_collisions) # further domain crossings are possible. if this happens, # velocity should be the same as original velocity b/c # immersed boundaries do not intersect with domain bndry, # so agent has slid off with original velocity heading. else: raise NameError ### Right boundary ### right_bool = np.logical_and(BC_bool, status_BC[:,dim]>0) right_idx = idx_array[right_bool] if len(right_idx) > 0: if bndry[1] == 'zero': # mask everything exiting on the right self.positions[right_idx, :] = ma.masked self.velocities[right_idx, :] = ma.masked self.accelerations[right_idx, :] = ma.masked # no further BC checks are made: masked entries are skipped elif bndry[1] == 'noflux': # agent slides along flat boundary. pos/vel/accel in dir of # boundary is zero. # additional IB crossings are possible, so first find # point of intersection with the boundary to enable this # check. if self.envir.ibmesh is not None and ib_collisions is not None: s_array = (self.positions[right_idx,dim]- self.envir.L[dim])/self.velocities[right_idx,dim] startpts = self.positions[right_idx,:] - (np.tile(s_array, (self.velocities.shape[1],1)).T* self.velocities[right_idx,:]) # now update pos/vel/accel self.positions[right_idx, dim] = self.envir.L[dim] self.velocities[right_idx, dim] = 0 self.accelerations[right_idx, dim] = 0 # now check for IB crossings. However, due to potential # complex interactions with the noflux boundary, enforce # sticky ib collisions in all cases. if self.envir.ibmesh is not None and ib_collisions is not None: for n, idx in enumerate(right_idx): startpt = startpts[n] endpt = self.positions[idx,:].copy() IBC_routine(idx, self, startpt, endpt, 'sticky') # further domain crossings remain possible elif bndry[1] == 'periodic': # wrap everything exiting on the right to the left self.positions[right_idx, dim] =\ np.mod(self.positions[right_idx, dim],self.envir.L[dim]) # check for IB crossings. first, get the point of re-entry # using the velocity. if self.envir.ibmesh is not None and ib_collisions is not None: s_array = (self.positions[right_idx,dim]-0)/ \ self.velocities[right_idx,dim] for n, idx in enumerate(right_idx): startpt = self.positions[idx,:] - \ s_array[n]*self.velocities[idx,:] endpt = self.positions[idx,:].copy() IBC_routine(idx, self, startpt, endpt, ib_collisions) # further domain crossings are possible. if this happens, # velocity should be the same as original velocity b/c # immersed boundaries do not intersect with domain bndry, # so agent has slid off with original velocity heading. else: raise NameError ##### All BC applied to first exit. Conduct recursion if necessary ##### if mult_idx is not None: self._domain_BC_loop(ib_collisions, idx_array=mult_idx) @staticmethod def _apply_internal_BC(startpt, endpt, mesh, max_meshpt_dist, old_intersection=None, kill=False, ib_collisions='sliding'): '''Apply internal boundaries to a trajectory starting and ending at startpt and endpt, returning a new endpt (or the original one) as appropriate. Parameters ---------- startpt : length 2 or 3 array start location for agent trajectory endpt : length 2 or 3 array end location for agent trajectory mesh : Nx2x2 or Nx3x3 array eligible mesh elements to check for intersection max_meshpt_dist : float max distance between two points on a mesh element (used to determine how far away from startpt to search for mesh elements) old_intersection : list-like of data (for internal use only) records the last intersection in the recursion to check if we are bouncing back and forth between two boundaries as a result of a concave angle and the right kind of trajectory vector. kill : bool (for internal use only) set to True in 3D case if we have previously slid along the boundary line between two mesh elements. This prevents such a thing from happening more than once, in case of pathological cases. ib_collisions : {'sliding' (default), 'sticky'} Type of interaction with immersed boundaries. In sliding collisions, conduct recursive vector projection until the length of the original vector is exhausted. In sticky collisions, just return the point of intersection. Returns ------- newendpt : length 2 or 3 array new end location for agent trajectory Acknowledgements --------------- Appreciation goes to Anne Ho, for pointing out that the centroid of an equalateral triangle is further away from any of its vertices than I had originally assumed it was. ''' if len(startpt) == 2: DIM = 2 else: DIM = 3 # In barycentric coordinates, the centoid is always (1/3,1/3,1/3), and # the entries of the coordinates must add to 1. This suggests that the # furthest away you can be from every vertex simultaneously is 2/3 # down the medians (an increase in one barycentric coordinate results # in a decrease in the others). Since the median length is bounded # above by the length of the longest side of the triangle, circles # centered at each vertex that are 2/3 times the length of the longest # triangle side should be sufficient to cover any triangle. # More precisely: equalateral triangles are probably the worst case. If # so, all medians have length l*sqrt(3/4), where l is the length of a # side of the triangle. This implies circles of radius l*sqrt(3/4)*2/3 # are a strict lower bound on covering any circle. # The result of this argument, from the worst-case scenario equalateral # triangle in our collection, forms the radius we need to search from # the line segment of travel in order to find vertices of mesh elements # that we potentially intersected. search_rad = max_meshpt_dist*2/3 seg_length_2 = np.linalg.norm(endpt - startpt)**2 def min_distance(pt_list): ''' Return minimum distances between line segment startpt,endpt and all points (rows) in pt_list. ''' if seg_length_2 == 0: return np.linalg.norm(pt_list-startpt,axis=1) # Consider the line extending the segment: startpt + t*(endpt-startpt) # Find the projection of all points onto this line. # It falls where t = [(p-startpt).(endpt-startpt)]/|startpt-endpt|**2 # We then clamp t from [0,1] to handle points outside the segment # startpt,endpt t_list = np.maximum(0,np.minimum(1,np.dot( pt_list-startpt,endpt-startpt)/seg_length_2)) # determine the projection points on the segment proj_pt_list = startpt + np.outer(t_list,(endpt-startpt)) # return the distance btwn pt_list and projected points return np.linalg.norm(pt_list-proj_pt_list,axis=1) # Find all mesh elements that have vertex points within search_rad of # the trajectory segment. pt_bool = min_distance( mesh.reshape((mesh.shape[0]*mesh.shape[1],mesh.shape[2])))<=search_rad pt_bool = pt_bool.reshape((mesh.shape[0],mesh.shape[1])) close_mesh = mesh[np.any(pt_bool,axis=1)] # Get intersections if DIM == 2: intersection = swarm._seg_intersect_2D(startpt, endpt, close_mesh[:,0,:], close_mesh[:,1,:]) else: intersection = swarm._seg_intersect_3D_triangles(startpt, endpt, close_mesh[:,0,:], close_mesh[:,1,:], close_mesh[:,2,:]) # Return endpt we already have if None. if intersection is None: return endpt # If we do have an intersection: if ib_collisions == 'sliding': # Project remaining piece of vector onto mesh and repeat processes # as necessary until we have a final result. return swarm._project_and_slide(startpt, endpt, intersection, close_mesh, max_meshpt_dist, DIM, old_intersection, kill) elif ib_collisions == 'sticky': # Return the point of intersection # small number to perturb off of the actual boundary in order to avoid # roundoff errors that would allow penetration EPS = 1e-7 back_vec = (startpt-endpt)/np.linalg.norm(endpt-startpt) return intersection[0] + back_vec*EPS @staticmethod def _project_and_slide(startpt, endpt, intersection, mesh, max_meshpt_dist, DIM, old_intersection=None, kill=False): '''Once we have an intersection point with an immersed mesh, project and slide the agent along the mesh for its remaining movement, and determine what happens if we fall off the edge of the element in all angle cases (e.g. some kind of recursion of detecting further intersections and resulting in additional vector projection) This, along with the intersection detection routines, is the real workhorse of immersed mesh interaction! Parameters ---------- startpt : length 2 or 3 array original start point of movement, before intersection endpt : length 2 or 3 array original end point of movement, w/o intersection intersection : list-like of data result of _seg_intersect_2D or _seg_intersect_3D_triangles. various information about the intersection with the immersed mesh element mesh : Nx2x2 or Nx3x3 array eligible (nearby) mesh elements for interaction max_meshpt_dist : float max distance between two points on a mesh element (used to determine how far away from startpt to search for mesh elements). Used here for passthrough to possible recursion DIM : int dimension of system, either 2 or 3 old_intersection : list-like of data (for internal use only) records the last intersection in the recursion to check if we are bouncing back and forth between two boundaries as a result of a concave angle and the right kind of trajectory vector. kill : bool (for internal use only) set to True in 3D case if we have previously slid along the boundary line between two mesh elements. This prevents such a thing from happening more than once, in case of pathological cases. Returns ------- newendpt : length 2 or 3 array new endpoint for movement after projection ''' # small number to perturb off of the actual boundary in order to avoid # roundoff errors that would allow penetration EPS = 1e-7 # Project remaining piece of vector from intersection onto mesh and get # a unit normal pointing out from the simplex # The first two entries of an intersection object are always: # 0) The point of intersection # 1) The fraction of line segment traveled before intersection occurred # NOTE: In the case of 2D, intersection[2] is a unit vector on the line # In the case of 3D, intersection[2] is a unit normal to the triangle # The remaining entries give the vertices of the line/triangle # Get leftover portion of travel vector vec = (1-intersection[1])*(endpt-startpt) if DIM == 2: # project vec onto the line proj = np.dot(vec,intersection[2])*intersection[2] # vec - proj is a normal that points from outside bndry inside # reverse direction and normalize to get unit vector pointing out norm_out_u = (proj-vec)/np.linalg.norm(proj-vec) if DIM == 3: # get a unit normal vector pointing back the way we came norm_out_u = -np.sign(np.dot(vec,intersection[2]))*intersection[2] # Get the component of vec that lies in the plane. We do this by # subtracting off the component which is in the normal direction proj = vec - np.dot(vec,norm_out_u)*norm_out_u # get a unit vector of proj for adding EPS in proj direction if np.linalg.norm(proj) > 0: proj_u = proj/np.linalg.norm(proj) else: proj_u = 0 # IMPORTANT: there can be roundoff error, so proj should be considered # an approximation to an in-boundary slide. # For this reason, pull newendpt back a bit for numerical stability newendpt = intersection[0] + proj + EPS*norm_out_u ################################ ########### 2D ########### ################################ if DIM == 2: # Detect sliding off 1D edge # Equivalent to going past the endpoints mesh_el_len = np.linalg.norm(intersection[4] - intersection[3]) Q0_dist = np.linalg.norm(newendpt-EPS*norm_out_u - intersection[3]) Q1_dist = np.linalg.norm(newendpt-EPS*norm_out_u - intersection[4]) # Since we are sliding on the mesh element, if the distance from # our new location to either of the mesh endpoints is greater # than the length of the mesh element, we must have gone beyond # the segment. if Q0_dist > mesh_el_len+EPS or Q1_dist > mesh_el_len+EPS: ########## Went past either Q0 or Q1 ########## # We check assuming the agent slid an additional EPS, because # if we end up in the non-concave case and fall off, we will # want to go EPS further so as to avoid corner penetration # The result of adding EPS in both the projection and normal # directions (below) is that some concave intersections will be # discovered (all acute, some obtuse) but others won't. However, # All undiscovered ones will be obtuse, so while re-detecting # the subsequent intersection is not the most efficient thing, # it should be stable. # check for new crossing of segment attached to the # current line segment # First, find adjacent mesh segments. These are ones that share # one of, but not both, endpoints with the current segment. # This will also pick up our current segment, but we shouldn't # be intersecting it. And if by some numerical error we do, # we need to treat it. pt_bool = np.logical_or( np.isclose(np.linalg.norm(mesh.reshape( (mesh.shape[0]*mesh.shape[1],mesh.shape[2]))-intersection[3], axis=1),0), np.isclose(np.linalg.norm(mesh.reshape( (mesh.shape[0]*mesh.shape[1],mesh.shape[2]))-intersection[4], axis=1),0) ) pt_bool = pt_bool.reshape((mesh.shape[0],mesh.shape[1])) adj_mesh = mesh[np.any(pt_bool,axis=1)] # Check for intersection with these segemtns, but translate # start/end points back from the segment a bit for numerical stability # This has already been done for newendpt # Also go EPS further than newendpt for stability for what follows if len(adj_mesh) > 0: adj_intersect = swarm._seg_intersect_2D(intersection[0]+EPS*norm_out_u, newendpt+EPS*proj_u, adj_mesh[:,0,:], adj_mesh[:,1,:]) else: adj_intersect = None if adj_intersect is not None: ########## Went past and intersected adjoining element! ########## # check that we haven't intersected this before # (should be impossible) if old_intersection is not None and\ np.all(adj_intersect[3] == old_intersection[3]) and\ np.all(adj_intersect[4] == old_intersection[4]): # Going back and forth! Movement stops here. # NOTE: This happens when 1) trying to go through a # mesh element, you 2) slide and intersect another # mesh element. The angle between elements is acute, # so you 3) slide back into the same element you # intersected in (1). # Back away from the intersection point slightly for # numerical stability and stay put. return adj_intersect[0] - EPS*proj_u # Otherwise: if Q0_dist > Q1_dist: # went past Q1; base vectors off Q1 vertex idx = 4 nidx = 3 else: # went past Q0; base vectors off Q0 vertex idx = 3 nidx = 4 vec0 = intersection[nidx] - intersection[idx] adj_idx = np.argmin([ np.linalg.norm(adj_intersect[3]-intersection[idx]), np.linalg.norm(adj_intersect[4]-intersection[idx])]) + 3 vec1 = adj_intersect[adj_idx] - intersection[idx] # Determine if the angle of mesh elements is acute or obtuse. if np.dot(vec0,vec1) >= 0: # Acute. Movement stops here. # Back away from the intersection point slightly for # numerical stability and stay put return adj_intersect[0] - EPS*proj_u else: # Obtuse. Repeat project_and_slide on new segment, # but send along info about old segment so we don't # get in an infinite loop. return swarm._project_and_slide(intersection[0]+EPS*norm_out_u, newendpt+EPS*proj_u, adj_intersect, mesh, max_meshpt_dist, DIM, intersection) ###### Went past, but did not intersect adjoining element! ###### # NOTE: There could still be an adjoining element at >180 degrees # But we will project along original heading as if there isn't one # and subsequently detect any intersections. # DIM == 2, adj_intersect is None if Q0_dist > Q1_dist: ##### went past Q1 ##### # put a new start point at the point crossing+EPS and bring out # EPS*norm_out_u newstartpt = intersection[4] + EPS*proj_u + EPS*norm_out_u elif Q0_dist < Q1_dist: ##### went past Q0 ##### newstartpt = intersection[3] + EPS*proj_u + EPS*norm_out_u else: raise RuntimeError("Impossible case??") # continue full amount of remaining movement along original heading # starting at newstartpt orig_unit_vec = (endpt-startpt)/np.linalg.norm(endpt-startpt) newendpt = newstartpt + np.linalg.norm(newendpt-newstartpt)*orig_unit_vec # repeat process to look for additional intersections # pass along current intersection in case of obtuse concave case return swarm._apply_internal_BC(newstartpt, newendpt, mesh, max_meshpt_dist, intersection) else: ########## Did not go past either Q0 or Q1 ########## # We simply end on the mesh element return newendpt ################################ ########### 3D ########### ################################ if DIM == 3: # Detect sliding off 2D edge using _seg_intersect_2D Q0_list = np.array(intersection[3:]) Q1_list = Q0_list[(1,2,0),:] # go a little further along trajectory to treat acute case tri_intersect = swarm._seg_intersect_2D(intersection[0] + EPS*norm_out_u, newendpt + EPS*proj_u, Q0_list, Q1_list) # if we reach the triangle boundary, check for a acute crossing # then project overshoot onto original heading and add to # intersection point if tri_intersect is not None: ########## Went past triangle boundary ########## # check for new crossing of simplex attached to the current triangle # First, find adjacent mesh segments. These are ones that share # one of, but not both, endpoints with the current segment. # This will also pick up our current segment, but we shouldn't # be intersecting it. And if by some numerical error we do, # we need to treat it. pt_bool = np.logical_or(np.logical_or( np.isclose(np.linalg.norm(mesh.reshape( (mesh.shape[0]*mesh.shape[1],mesh.shape[2]))-intersection[3], axis=1),0), np.isclose(np.linalg.norm(mesh.reshape( (mesh.shape[0]*mesh.shape[1],mesh.shape[2]))-intersection[4], axis=1),0)), np.isclose(np.linalg.norm(mesh.reshape( (mesh.shape[0]*mesh.shape[1],mesh.shape[2]))-intersection[5], axis=1),0) ) pt_bool = pt_bool.reshape((mesh.shape[0],mesh.shape[1])) adj_mesh = mesh[np.any(pt_bool,axis=1)] # check for intersection, but translate start/end points back # from the simplex a bit for numerical stability # this has already been done for newendpt # also go EPS further than newendpt for stability for what follows if len(adj_mesh) > 0: adj_intersect = swarm._seg_intersect_3D_triangles( intersection[0]+EPS*norm_out_u, newendpt+EPS*proj_u, adj_mesh[:,0,:], adj_mesh[:,1,:], adj_mesh[:,2,:]) else: adj_intersect = None if adj_intersect is not None: ########## Intersected adjoining element! ########## # Acute or slightly obtuse intersection. In 3D, we will slide # regardless of if it is acute or not. # First, detect if we're hitting a simplex we've already # been on before. If so, we follow the intersection line. if old_intersection is not None and\ np.all(adj_intersect[3] == old_intersection[3]) and\ np.all(adj_intersect[4] == old_intersection[4]) and\ np.all(adj_intersect[5] == old_intersection[5]): # Going back and forth between two elements. Project # onto the line of intersection. # NOTE: This happens when 1) trying to go through a # mesh element, you 2) slide and intersect another # mesh element. The angle between elements is acute, # so you 3) slide back into the same element you # intersected in (1). # NOTE: In rare cases, solving this problem by following # the line of intersection can result in an infinite # loop. For example, when going toward the point of # a tetrahedron. This is what the kill switch is for. if kill: # End here to prevent possible bad behavior. return adj_intersect[0] - EPS*proj_u kill = True # Find the points in common between adj_intersect and # intersection # Enter debugging here to test algorithm # import pdb; pdb.set_trace() dist_mat = distance.cdist(intersection[3:],adj_intersect[3:]) pt_bool = np.isclose(dist_mat, np.zeros_like(dist_mat), atol=1e-6).any(1) two_pts = np.array(intersection[3:])[pt_bool] assert two_pts.shape[0] == 2, "Other than two points found?" assert two_pts.shape[1] == 3, "Points aren't 3D?" # Get a unit vector from these points vec_intersect = two_pts[1,:] - two_pts[0,:] vec_intersect /= np.linalg.norm(vec_intersect) # Back up a tad from the intersection point, and project # leftover movement in the direction of intersection vector newstartpt = adj_intersect[0] - EPS*proj_u # Get leftover portion of the travel vector vec_to_project = (1-adj_intersect[1])*proj # project vec onto the line proj_vec = np.dot(vec_to_project,vec_intersect)*vec_intersect # Get new endpoint newendpt = newstartpt + proj_vec # Check for more intersections return swarm._apply_internal_BC(newstartpt, newendpt, mesh, max_meshpt_dist, adj_intersect, kill) # Not an already discovered mesh element. # We slide. Pass along info about the old element. return swarm._project_and_slide(intersection[0]+EPS*norm_out_u, newendpt+EPS*proj_u+EPS*norm_out_u, adj_intersect, mesh, max_meshpt_dist, DIM, intersection, kill) ########## Did not intersect adjoining element! ########## # NOTE: There could still be an adjoining element w/ >180 connection # But we will project along original heading as if there isn't # one and subsequently detect any intersections. # DIM == 3, adj_intersect is None # put the new start point on the edge of the simplex (+EPS) and # continue full amount of remaining movement along original heading newstartpt = tri_intersect[0] + EPS*proj_u + EPS*norm_out_u orig_unit_vec = (endpt-startpt)/np.linalg.norm(endpt-startpt) # norm(newendpt - tri_intersect[0]) is the length of the overshoot. newendpt = newstartpt + np.linalg.norm(newendpt-tri_intersect[0])*orig_unit_vec # repeat process to look for additional intersections return swarm._apply_internal_BC(newstartpt, newendpt, mesh, max_meshpt_dist, intersection, kill) else: # otherwise, we end on the mesh element return newendpt @staticmethod def _seg_intersect_2D(P0, P1, Q0_list, Q1_list, get_all=False): '''Find the intersection between two line segments (2D objects), P and Q, returning None if there isn't one or if they are parallel. If Q is a 2D array, loop over the rows of Q finding all intersections between P and each row of Q, but only return the closest intersection to P0 (if there is one, otherwise None) This works for both 2D problems and problems in which P is a 3D line segment roughly lying on a plane (e.g., in cases of projection along a 3D triangular mesh element). The plane is described by the first two vectors Q, so in this case, Q0_list and Q1_list must have at least two rows. The 3D problem is robust to cases where P is not exactly in the plane because the algorithm is actually checking to see if its projection onto the triangle crosses any of the lines in Q. This is important to deal with roundoff error. This algorithm uses a parameteric equation approach for speed, based on [1]_ Parameters ---------- P0 : length 2 (or 3) array first point in line segment P P1 : length 2 (or 3) array second point in line segment P Q0_list : Nx2 (Nx3) ndarray first points in a list of line segments. Q1_list : Nx2 (Nx3) ndarray second points in a list of line segments. get_all : bool Return all intersections instead of just the first one encountered as one travels from P0 to P1. Returns ------- None if there is no intersection. Otherwise: x : length 2 (or 3) array the coordinates of the point of first intersection s_I : float between 0 and 1 the fraction of the line segment traveled from P0 to P1 before intersection occurred vec : length 2 (or 3) array directional unit vector along the boundary (Q) intersected Q0 : length 2 (or 3) array first endpoint of mesh segment intersected Q1 : length 2 (or 3) array second endpoint of mesh segment intersected References ---------- .. [1] Sunday, Daniel, (2021). Practial Geometry Algorithms with C++ Code, self-published: Amazon KDP. ''' u = P1 - P0 v = Q1_list - Q0_list w = P0 - Q0_list if len(P0) == 2: u_perp = np.array([-u[1], u[0]]) v_perp = np.array([-v[...,1], v[...,0]]).T else: normal = np.cross(v[0],v[1]) normal /= np.linalg.norm(normal) # roundoff error in only an np.dot projection can be as high as 1e-7 assert np.isclose(np.dot(u,normal),0,atol=1e-6), "P vector not parallel to Q plane" u_perp = np.cross(u,normal) v_perp = np.cross(v,normal) if len(Q0_list.shape) == 1: # only one point in Q list denom = np.dot(v_perp,u) if denom != 0: s_I = np.dot(-v_perp,w)/denom t_I = -np.dot(u_perp,w)/denom if 0<=s_I<=1 and 0<=t_I<=1: return (P0 + s_I*u, s_I, v, Q0_list, Q1_list) return None denom_list = np.multiply(v_perp,u).sum(1) #vectorized dot product # We need to deal with parallel cases. With roundoff error, exact zeros # are unlikely (but ruled out below). Another possiblity is getting # inf values, but these will not record as being between 0 and 1, and # will not throw errors when compared to these values. So all should # be good. # # Non-parallel cases: not_par = denom_list != 0 # All non-parallel lines in the same plane intersect at some point. # Find the parametric values for both vectors at which that intersect # happens. Call these s_I and t_I respectively. s_I_list = -np.ones_like(denom_list) t_I_list = -np.ones_like(denom_list) # Now only need to calculuate s_I & t_I for non parallel cases; others # will report as not intersecting automatically (as -1) # (einsum is faster for vectorized dot product, but need same length, # non-empty vectors) # In 3D, we are actually projecting u onto the plane of the triangle # and doing our calculation there. So no problem about numerical # error and offsets putting u above the plane. if np.any(not_par): s_I_list[not_par] = np.einsum('ij,ij->i',-v_perp[not_par],w[not_par])/denom_list[not_par] t_I_list[not_par] = -np.multiply(u_perp,w[not_par]).sum(1)/denom_list[not_par] # The length of our vectors parameterizing the lines is the same as the # length of the line segment. So for the line segments to have intersected, # the parameter values for their intersect must both be in the unit interval. intersect = np.logical_and( np.logical_and(0<=s_I_list, s_I_list<=1), np.logical_and(0<=t_I_list, t_I_list<=1)) if np.any(intersect): if get_all: x = [] for s_I in s_I_list[intersect]: x.append(P0+s_I*u) return zip( x, s_I_list[intersect], v[intersect]/np.linalg.norm(v[intersect]), Q0_list[intersect], Q1_list[intersect] ) else: # find the closest intersection and return it Q0 = Q0_list[intersect] Q1 = Q1_list[intersect] v_intersected = v[intersect] s_I = s_I_list[intersect].min() s_I_idx = s_I_list[intersect].argmin() return (P0 + s_I*u, s_I, v_intersected[s_I_idx]/np.linalg.norm(v_intersected[s_I_idx]), Q0[s_I_idx], Q1[s_I_idx]) else: return None @staticmethod def _seg_intersect_3D_triangles(P0, P1, Q0_list, Q1_list, Q2_list, get_all=False): '''Find the intersection between a line segment P0 to P1 and any of the triangles given by Q0, Q1, Q2 where each row across the three arrays is a different triangle (three points). Returns None if there is no intersection. This algorithm uses a parameteric equation approach for speed, based on [2]_ Parameters ---------- P0 : length 3 array first point in line segment P P1 : length 3 array second point in line segment P Q0 : Nx3 ndarray first points in a list of triangles. Q1 : Nx3 ndarray second points in a list of triangles. Q2 : Nx3 ndarray third points in a list of triangles. get_all : bool Return all intersections instead of just the first one encountered as you travel from P0 to P1. Returns ------- None if there is no intersection. Otherwise: x : length 3 array the coordinates of the first point of intersection s_I : float between 0 and 1 the fraction of the line segment traveled from P0 before intersection occurred (only if intersection occurred) normal : length 3 array normal unit vector to plane of intersection Q0 : length 3 array first vertex of triangle intersected Q1 : length 3 array second vertex of triangle intersected Q2 : length 3 array third vertex of triangle intersected References ---------- .. [1] Sunday, Daniel, (2021). Practial Geometry Algorithms with C++ Code, self-published: Amazon KDP. ''' # First, determine the intersection between the line and the plane # Get normal vectors Q1Q0_diff = Q1_list-Q0_list Q2Q0_diff = Q2_list-Q0_list n_list = np.cross(Q1Q0_diff, Q2Q0_diff) u = P1 - P0 w = P0 - Q0_list # At intersection, w + su is perpendicular to n if len(Q0_list.shape) == 1: # only one triangle denom = np.dot(n_list,u) if denom != 0: s_I = np.dot(-n_list,w)/denom if 0<=s_I<=1: # line segment crosses full plane cross_pt = P0 + s_I*u # calculate barycentric coordinates normal = n_list/np.linalg.norm(n_list) A_dbl = np.dot(n_list, normal) Q0Pt = cross_pt-Q0_list A_u_dbl = np.dot(np.cross(Q0Pt,Q2Q0_diff),normal) A_v_dbl = np.dot(np.cross(Q1Q0_diff,Q0Pt),normal) coords = np.array([A_u_dbl/A_dbl, A_v_dbl/A_dbl, 0]) coords[2] = 1 - coords[0] - coords[1] # check if point is in triangle if np.all(coords>=0): return (cross_pt, s_I, normal, Q0_list, Q1_list, Q2_list) return None denom_list = np.multiply(n_list,u).sum(1) #vectorized dot product # record non-parallel cases not_par = denom_list != 0 # default is not intersecting s_I_list = -np.ones_like(denom_list) # get intersection parameters # (einsum is faster for vectorized dot product, but need same length vectors) if np.any(not_par): s_I_list[not_par] = np.einsum('ij,ij->i',-n_list[not_par],w[not_par])/denom_list[not_par] # test for intersection of line segment with full plane plane_int = np.logical_and(0<=s_I_list, s_I_list<=1) # calculate barycentric coordinates for each plane intersection closest_int = (None, -1, None) intersections = [] for n, s_I in zip(np.arange(len(plane_int))[plane_int], s_I_list[plane_int]): # if get_all is False, we only care about the closest triangle intersection! # see if we need to worry about this one if closest_int[1] == -1 or closest_int[1] > s_I or get_all: cross_pt = P0 + s_I*u normal = n_list[n]/np.linalg.norm(n_list[n]) A_dbl = np.dot(n_list[n], normal) Q0Pt = cross_pt-Q0_list[n] A_u_dbl = np.dot(np.cross(Q0Pt,Q2Q0_diff[n]),normal) A_v_dbl = np.dot(np.cross(Q1Q0_diff[n],Q0Pt),normal) coords = np.array([A_u_dbl/A_dbl, A_v_dbl/A_dbl, 0]) coords[2] = 1 - coords[0] - coords[1] # check if point is in triangle if np.all(coords>=0): if get_all: intersections.append((cross_pt, s_I, normal, Q0_list[n], Q1_list[n], Q2_list[n])) else: closest_int = (cross_pt, s_I, normal, Q0_list[n], Q1_list[n], Q2_list[n]) if not get_all: if closest_int[0] is None: return None else: return closest_int else: if len(intersections) == 0: return None else: return intersections @staticmethod def _dist_point_to_plane(P0, normal, Q0): '''Return the distance from the point P0 to the plane given by a normal vector and a point on the plane, Q0. For debugging.''' d = np.dot(normal, Q0) return np.abs(np.dot(normal,P0)-d)/np.linalg.norm(normal) def _calc_basic_stats(self, DIM3, t_indx=None): ''' Return basic stats about % agents remaining, fluid velocity, and agent velocity for plot printing. Parameters ---------- DIM3 : bool indicates the dimension of the domain (True for 3D) t_indx : int, optional The time index for pos_history or None for current time Returns ------- perc_left : float percentage of agents left within the domain avg_spd : float average fluid speed max_spd : float maximum fluid speed avg_spd_x : float average x-component of fluid velocity avg_spd_y : float average y-component of fluid velocity avg_spd_z : float, 3D only average z-component of fluid velocity avg_swrm_vel : array average agent velocity ''' # get % of agents left in domain if t_indx is None: num_left = self.positions[:,0].compressed().size else: num_left = self.pos_history[t_indx][:,0].compressed().size if len(self.pos_history) > 0: num_orig = self.pos_history[0][:,0].compressed().size else: num_orig = num_left perc_left = 100*num_left/num_orig # get average swarm velocity if t_indx is None: if np.all(self.velocities.mask): vel_data = np.zeros(self.velocities.shape) elif np.any(self.velocities.mask): vel_data = self.velocities[~self.velocities.mask.any(axis=1)] else: vel_data = self.velocities avg_swrm_vel = vel_data.mean(axis=0) elif t_indx == 0: avg_swrm_vel = np.zeros(len(self.envir.L)) else: vel_data = (self.pos_history[t_indx] - self.pos_history[t_indx-1])/( self.envir.time_history[t_indx]-self.envir.time_history[t_indx-1]) avg_swrm_vel = vel_data.mean(axis=0) if self.envir.flow is None and not DIM3: return perc_left, 0, 0, 0, 0, avg_swrm_vel elif self.envir.flow is None and DIM3: return perc_left, 0, 0, 0, 0, 0, avg_swrm_vel if not DIM3: # 2D flow # get current fluid flow info if len(self.envir.flow[0].shape) == 2: # temporally constant flow flow = self.envir.flow else: # temporally changing flow flow = self.envir.interpolate_temporal_flow(t_indx=t_indx) flow_spd = np.sqrt(flow[0]**2 + flow[1]**2) avg_spd_x = flow[0].mean() avg_spd_y = flow[1].mean() avg_spd = flow_spd.mean() max_spd = flow_spd.max() return perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_swrm_vel else: # 3D flow if len(self.envir.flow[0].shape) == 3: # temporally constant flow flow = self.envir.flow else: # temporally changing flow flow = self.envir.interpolate_temporal_flow(t_indx) flow_spd = np.sqrt(flow[0]**2 + flow[1]**2 + flow[2]**2) avg_spd_x = flow[0].mean() avg_spd_y = flow[1].mean() avg_spd_z = flow[2].mean() avg_spd = flow_spd.mean() max_spd = flow_spd.max() return perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_spd_z, avg_swrm_vel
[docs] def plot(self, t=None, filename=None, blocking=True, dist='density', fluid=None, clip=None, figsize=None, save_kwargs=None, azim=None, elev=None): '''Plot the position of the swarm at time t, or at the current time if no time is supplied. The actual time plotted will depend on the history of movement steps; the closest entry in environment.time_history will be shown without interpolation. Parameters ---------- t : float, optional time to plot. if None (default), the current time. filename : str, optional file name to save image as. Image will not be shown, only saved. blocking : bool, default True whether the plot should block execution or not dist : {'density' (default), 'cov', float, 'hist'} whether to plot Gaussian kernel density estimation or histogram. Options are: * 'density': plot Gaussian KDE using Scotts Factor from scipy.stats.gaussian_kde * 'cov': use the variance in each direction from self.shared_props['cov'] to plot Gaussian KDE * float: plot Gaussian KDE using the given bandwidth factor to multiply the KDE variance by * 'hist': plot histogram fluid : {'vort', 'quiver'}, optional Plot info on the fluid in the background. 2D only! If None, don't plot anything related to the fluid. Options are: * 'vort': plot vorticity in the background * 'quiver': quiver plot of fluid velocity in the background clip : float, optional if plotting vorticity, specifies the clip value for pseudocolor. this value is used for both negative and positive vorticity. figsize : tuple of length 2, optional figure size in inches, (width, height). default is a heurstic that works... most of the time? save_kwargs : dict of keyword arguments, optional keys must be valid strings that match keyword arguments for the matplotlib savefig function. These arguments will be passed to savefig assuming that a filename has been specified. azim : float, optional In 3D plots, the azimuthal viewing angle. Defaults to -60. elev : float, optional In 3D plots, the elevation viewing angle. Defaults to 30. ''' if t is not None and len(self.envir.time_history) != 0: loc = np.searchsorted(self.envir.time_history, t) if loc == len(self.envir.time_history): if (t-self.envir.time_history[-1]) > (self.envir.time-t): loc = None else: loc = -1 elif t < self.envir.time_history[loc]: if (self.envir.time_history[loc]-t) > (t-self.envir.time_history[loc-1]): loc -= 1 else: loc = None # get time and positions if loc is None: time = self.envir.time positions = self.positions else: time = self.envir.time_history[loc] positions = self.pos_history[loc] if len(self.envir.L) == 2: # 2D plot if figsize is None: aspectratio = self.envir.L[0]/self.envir.L[1] if aspectratio > 1: x_length = np.min((6*aspectratio,12)) y_length = 6 elif aspectratio < 1: x_length = 6 y_length = np.min((6/aspectratio,8)) else: x_length = 6 y_length = 6 fig = plt.figure(figsize=(x_length,y_length)) else: fig = plt.figure(figsize=figsize) ax, axHistx, axHisty = self.envir._plot_setup(fig) if figsize is None: # some final adjustments in a particular case if x_length == 12: ax_pos = ax.get_position().get_points() axHx_pos = np.array(axHistx.get_position().get_points()) axHy_pos = np.array(axHisty.get_position().get_points()) if ax_pos[0,1] > 0.1: extra = 2*(ax_pos[0,1] - 0.1)*y_length fig.set_size_inches(x_length,y_length-extra) prop = (y_length-extra/4)/y_length prop_wdth = (y_length-extra/2)/y_length prop_len = (y_length-extra)/y_length axHistx.set_position([axHx_pos[0,0],axHx_pos[0,1]*prop, axHx_pos[1,0]-axHx_pos[0,0], (axHx_pos[1,1]-axHx_pos[0,1])/prop_wdth]) axHisty.set_position([axHy_pos[0,0],axHy_pos[0,1]*prop_len, axHy_pos[1,0]-axHy_pos[0,0], (axHy_pos[1,1]-axHy_pos[0,1])/prop_len]) # fluid visualization if fluid == 'vort' and self.envir.flow is not None: vort = self.envir.get_2D_vorticity(t_indx=loc) if clip is not None: norm = colors.Normalize(-abs(clip),abs(clip),clip=True) else: norm = None ax.pcolormesh(self.envir.flow_points[0], self.envir.flow_points[1], vort.T, shading='gouraud', cmap='RdBu', norm=norm, alpha=0.9, antialiased=True) elif fluid == 'quiver' and self.envir.flow is not None: # get dimensions of axis to estimate a decent quiver density ax_pos = ax.get_position().get_points() fig_size = fig.get_size_inches() wdth_inch = fig_size[0]*(ax_pos[1,0]-ax_pos[0,0]) height_inch = fig_size[1]*(ax_pos[1,1]-ax_pos[0,1]) # use about 4.15/inch density of arrows x_num = round(4.15*wdth_inch) y_num = round(4.15*height_inch) M = int(round(len(self.envir.flow_points[0])/x_num)) N = int(round(len(self.envir.flow_points[1])/y_num)) # get worse case max velocity vector for scaling max_u = self.envir.flow[0].max(); max_v = self.envir.flow[1].max() max_mag = np.linalg.norm(np.array([max_u,max_v])) if len(self.envir.flow[0].shape) > 2: flow = self.envir.interpolate_temporal_flow(t_indx=loc) else: flow = self.envir.flow ax.quiver(self.envir.flow_points[0][::M], self.envir.flow_points[1][::N], flow[0][::M,::N].T, flow[1][::M,::N].T, scale=max_mag*5, alpha=0.2) # scatter plot and time text ax.scatter(positions[:,0], positions[:,1], label='organism', c='darkgreen', s=3) ax.text(0.02, 0.95, 'time = {:.2f}'.format(time), transform=ax.transAxes, fontsize=12) # textual info perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_swrm_vel = \ self._calc_basic_stats(DIM3=False, t_indx=loc) plt.figtext(0.77, 0.77, '{:.1f}% remain\n'.format(perc_left)+ '\n ------ Info ------\n'+ r'Fluid $v_{max}$'+': {:.1g} {}/s\n'.format(max_spd, self.envir.units)+ r'Fluid $\overline{v}$'+': {:.1g} {}/s\n'.format(avg_spd, self.envir.units)+ r'Agent $\overline{v}$'+': {:.1g} {}/s\n'.format(np.linalg.norm(avg_swrm_vel), self.envir.units), fontsize=10) axHistx.text(0.01, 0.98, r'Fluid $\overline{v}_x$'+': {:.2g} \n'.format(avg_spd_x)+ r'Agent $\overline{v}_x$'+': {:.2g}'.format(avg_swrm_vel[0]), transform=axHistx.transAxes, verticalalignment='top', fontsize=10) axHisty.text(0.02, 0.99, r'Fluid $\overline{v}_y$'+': {:.2g} \n'.format(avg_spd_y)+ r'Agent $\overline{v}_y$'+': {:.2g}'.format(avg_swrm_vel[1]), transform=axHisty.transAxes, verticalalignment='top', fontsize=10) if dist == 'hist': # histograms bins_x = np.linspace(0, self.envir.L[0], 26) bins_y = np.linspace(0, self.envir.L[1], 26) axHistx.hist(positions[:,0].compressed(), bins=bins_x) axHisty.hist(positions[:,1].compressed(), bins=bins_y, orientation='horizontal') else: # Gaussian Kernel Density Estimation if dist == 'cov': fac_x = self.shared_props['cov'][0,0] fac_y = self.shared_props['cov'][1,1] else: try: fac_x = float(dist) fac_y = fac_x except: fac_x = None fac_y = None xmesh = np.linspace(0, self.envir.L[0], 1000) ymesh = np.linspace(0, self.envir.L[1], 1000) # deal with point sources pos_x = positions[:,0].compressed() pos_y = positions[:,1].compressed() try: if len(pos_x) > 1: x_density = stats.gaussian_kde(pos_x, fac_x) x_density = x_density(xmesh) elif len(pos_x) == 1: raise np.linalg.LinAlgError else: x_density = np.zeros_like(xmesh) except np.linalg.LinAlgError: idx = (np.abs(xmesh - pos_x[0])).argmin() x_density = np.zeros_like(xmesh); x_density[idx] = 1 try: if len(pos_y) > 1: y_density = stats.gaussian_kde(pos_y, fac_y) y_density = y_density(ymesh) elif len(pos_y) == 1: raise np.linalg.LinAlgError else: y_density = np.zeros_like(ymesh) except np.linalg.LinAlgError: idy = (np.abs(ymesh - pos_y[0])).argmin() y_density = np.zeros_like(ymesh); y_density[idy] = 1 axHistx.plot(xmesh, x_density) axHisty.plot(y_density, ymesh) axHistx.get_yaxis().set_ticks([]) axHisty.get_xaxis().set_ticks([]) if np.max(x_density) != 0: axHistx.set_ylim(bottom=0, top=np.max(x_density)) else: axHistx.set_ylim(bottom=0) if np.max(y_density) != 0: axHisty.set_xlim(left=0, right=np.max(y_density)) else: axHisty.set_xlim(left=0) else: # 3D plot if figsize is None: fig = plt.figure(figsize=(10,5)) else: fig = plt.figure(figsize=figsize) ax, axHistx, axHisty, axHistz = self.envir._plot_setup(fig) if azim is not None or elev is not None: ax.view_init(elev, azim) # scatter plot and time text ax.scatter(positions[:,0], positions[:,1], positions[:,2], label='organism') ax.text2D(0.02, 1, 'time = {:.2f}'.format(time), transform=ax.transAxes, verticalalignment='top', fontsize=12) # textual info perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_spd_z, avg_swrm_vel = \ self._calc_basic_stats(DIM3=True, t_indx=loc) ax.text2D(0.75, 0.9, r'Fluid $v_{max}$'+': {:.2g} {}/s\n'.format(max_spd, self.envir.units)+ r'Fluid $v_{avg}$'+': {:.2g} {}/s\n'.format(avg_spd, self.envir.units)+ r'Agent $v_{avg}$'+': {:.2g} {}/s'.format(np.linalg.norm(avg_swrm_vel), self.envir.units), transform=ax.transAxes, horizontalalignment='left', fontsize=10) ax.text2D(0.02, 0, '{:.1f}% remain\n'.format(perc_left), transform=fig.transFigure, fontsize=10) axHistx.text(0.02, 0.98, r'Fluid $\overline{v}_x$'+': {:.2g} {}/s\n'.format(avg_spd_x, self.envir.units)+ r'Agent $\overline{v}_x$'+': {:.2g} {}/s'.format(avg_swrm_vel[0], self.envir.units), transform=axHistx.transAxes, verticalalignment='top', fontsize=10) axHisty.text(0.02, 0.98, r'Fluid $\overline{v}_y$'+': {:.2g} {}/s\n'.format(avg_spd_y, self.envir.units)+ r'Agent $\overline{v}_y$'+': {:.2g} {}/s'.format(avg_swrm_vel[1], self.envir.units), transform=axHisty.transAxes, verticalalignment='top', fontsize=10) axHistz.text(0.02, 0.98, r'Fluid $\overline{v}_z$'+': {:.2g} {}/s\n'.format(avg_spd_z, self.envir.units)+ r'Agent $\overline{v}_z$'+': {:.2g} {}/s'.format(avg_swrm_vel[2], self.envir.units), transform=axHistz.transAxes, verticalalignment='top', fontsize=10) if dist == 'hist': # histograms bins_x = np.linspace(0, self.envir.L[0], 26) bins_y = np.linspace(0, self.envir.L[1], 26) bins_z = np.linspace(0, self.envir.L[2], 26) axHistx.hist(positions[:,0].compressed(), bins=bins_x, alpha=0.8) axHisty.hist(positions[:,1].compressed(), bins=bins_y, alpha=0.8) axHistz.hist(positions[:,2].compressed(), bins=bins_z, alpha=0.8) else: # Gaussian Kernel Density Estimation if dist == 'cov': fac_x = self.shared_props['cov'][0,0] fac_y = self.shared_props['cov'][1,1] fac_z = self.shared_props['cov'][2,2] else: try: fac_x = float(dist) fac_y = fac_x fac_z = fac_x except: fac_x = None fac_y = None fac_z = None xmesh = np.linspace(0, self.envir.L[0], 1000) ymesh = np.linspace(0, self.envir.L[1], 1000) zmesh = np.linspace(0, self.envir.L[2], 1000) # deal with point sources pos_x = positions[:,0].compressed() pos_y = positions[:,1].compressed() pos_z = positions[:,2].compressed() try: if len(pos_x) > 1: x_density = stats.gaussian_kde(pos_x, fac_x) x_density = x_density(xmesh) elif len(pos_x) == 1: raise np.linalg.LinAlgError else: x_density = np.zeros_like(xmesh) except np.linalg.LinAlgError: idx = (np.abs(xmesh - pos_x[0])).argmin() x_density = np.zeros_like(xmesh); x_density[idx] = 1 try: if len(pos_y) > 1: y_density = stats.gaussian_kde(pos_y, fac_y) y_density = y_density(ymesh) elif len(pos_y) == 1: raise np.linalg.LinAlgError else: y_density = np.zeros_like(ymesh) except np.linalg.LinAlgError: idy = (np.abs(ymesh - pos_y[0])).argmin() y_density = np.zeros_like(ymesh); y_density[idy] = 1 try: if len(pos_z) > 1: z_density = stats.gaussian_kde(pos_z, fac_z) z_density = z_density(zmesh) elif len(pos_z) == 1: raise np.linalg.LinAlgError else: z_density = np.zeros_like(zmesh) except np.linalg.LinAlgError: idz = (np.abs(zmesh - pos_z[0])).argmin() z_density = np.zeros_like(zmesh); z_density[idz] = 1 axHistx.plot(xmesh, x_density) axHisty.plot(ymesh, y_density) axHistz.plot(zmesh, z_density) axHistx.get_yaxis().set_ticks([]) axHisty.get_yaxis().set_ticks([]) axHistz.get_yaxis().set_ticks([]) if np.max(x_density) != 0: axHistx.set_ylim(bottom=0, top=np.max(x_density)) else: axHistx.set_ylim(bottom=0) if np.max(y_density) != 0: axHisty.set_ylim(bottom=0, top=np.max(y_density)) else: axHisty.set_ylim(bottom=0) if np.max(z_density) != 0: axHistz.set_ylim(bottom=0, top=np.max(z_density)) else: axHistz.set_ylim(bottom=0) # show the plot if filename is None: plt.show(block=blocking) else: if save_kwargs is not None: plt.savefig(filename, **save_kwargs) else: plt.savefig(filename)
[docs] def plot_all(self, movie_filename=None, frames=None, downsamp=None, fps=10, dist='density', fluid=None, clip=None, figsize=None, save_kwargs=None, writer_kwargs=None, azim=None, elev=None): ''' Plot the history of the swarm's movement, incl. current time in successively updating plots or saved as a movie file. A movie file is created if movie_filename is specified. Parameters ---------- movie_filename : string, optional file name to save movie as. file extension will determine the type of file saved. frames : iterable of integers, optional. If None, plot the entire history of the swarm's movement including the present time, with each step being a frame in the animation. If an iterable, plot only the time steps of the swarm as indexed by the iterable (note, this is an interable of the time step indices, not the time in seconds at those time steps!). downsamp : iterable of int or int, optional If None, do not downsample the agents - plot them all. If an integer, plot only the first n agents (equivalent to range(downsamp)). If an iterable, plot only the agents specified. In all cases, statistics are reported for the TOTAL population, both shown and unshown. This includes the histograms/KDE plots. fps : int, default=10 Frames per second, only used if saving a movie to file. Make sure this is at least as big as 1/dt, where dt is the time interval between frames! dist : {'density' (default), 'cov', float, 'hist'} whether to plot Gaussian kernel density estimation or histogram. Options are: * 'density': plot Gaussian KDE using Scotts Factor from scipy.stats.gaussian_kde * 'cov': use the variance in each direction from self.shared_props['cov'] to plot Gaussian KDE * float: plot Gaussian KDE using the given bandwidth factor to multiply the KDE variance by * 'hist': plot histogram fluid : {'vort', 'quiver'}, optional Plot info on the fluid in the background. 2D only! If None, don't plot anything related to the fluid. Options are: * 'vort': plot vorticity in the background * 'quiver': quiver plot of fluid velocity in the background clip : float, optional if plotting vorticity, specifies the clip value for pseudocolor. this value is used for both negative and positive vorticity. figsize : tuple of length 2, optional figure size in inches, (width, height). default is a heurstic that works... most of the time? writer_kwargs : dict of keyword arguments, optional keys must be valid strings that match keyword arguments for a matplotlib save_kwargs : dict of keyword arguments, optional keys must be valid strings that match keyword arguments for the matplotlib animation.FFMpegWriter object. These arguments will be used in the writer object initiation save assuming that a movie_filename has been specified. Otherwise, defaults are the passed in fps and metadata=dict(artist='Christopher Strickland')). azim : float, optional In 3D plots, the azimuthal viewing angle. Defaults to -60. elev : float, optional In 3D plots, the elevation viewing angle. Defaults to 30. ''' if len(self.envir.time_history) == 0: print('No position history! Plotting current position...') self.plot() return if movie_filename is not None: print("Creating video... this could take a long time!") DIM3 = (len(self.envir.L) == 3) if frames is None: n0 = 0 else: n0 = frames[0] if isinstance(downsamp, int): downsamp = range(downsamp) if not DIM3: ### 2D setup ### if figsize is None: aspectratio = self.envir.L[0]/self.envir.L[1] if aspectratio > 1: x_length = np.min((6*aspectratio,12)) y_length = 6 elif aspectratio < 1: x_length = 6 y_length = np.min((6/aspectratio,8)) else: x_length = 6 y_length = 6 fig = plt.figure(figsize=(x_length,y_length)) else: fig = plt.figure(figsize=figsize) ax, axHistx, axHisty = self.envir._plot_setup(fig) if figsize is None: # some final adjustments in a particular case if x_length == 12: ax_pos = ax.get_position().get_points() axHx_pos = np.array(axHistx.get_position().get_points()) axHy_pos = np.array(axHisty.get_position().get_points()) if ax_pos[0,1] > 0.1: extra = 2*(ax_pos[0,1] - 0.1)*y_length fig.set_size_inches(x_length,y_length-extra) prop = (y_length-extra/4)/y_length prop_wdth = (y_length-extra/2)/y_length prop_len = (y_length-extra)/y_length axHistx.set_position([axHx_pos[0,0],axHx_pos[0,1]*prop, axHx_pos[1,0]-axHx_pos[0,0], (axHx_pos[1,1]-axHx_pos[0,1])/prop_wdth]) axHisty.set_position([axHy_pos[0,0],axHy_pos[0,1]*prop_len, axHy_pos[1,0]-axHy_pos[0,0], (axHy_pos[1,1]-axHy_pos[0,1])/prop_len]) # fluid visualization if fluid == 'vort' and self.envir.flow is not None: if clip is not None: norm = colors.Normalize(-abs(clip),abs(clip),clip=True) else: norm = None fld = ax.pcolormesh([self.envir.flow_points[0]], self.envir.flow_points[1], np.zeros(self.envir.flow[0].shape[1:]).T, shading='gouraud', cmap='RdBu', norm=norm, alpha=0.9) elif fluid == 'quiver' and self.envir.flow is not None: # get dimensions of axis to estimate a decent quiver density ax_pos = ax.get_position().get_points() fig_size = fig.get_size_inches() wdth_inch = fig_size[0]*(ax_pos[1,0]-ax_pos[0,0]) height_inch = fig_size[1]*(ax_pos[1,1]-ax_pos[0,1]) # use about 4.15/inch density of arrows x_num = round(4.15*wdth_inch) y_num = round(4.15*height_inch) M = round(len(self.envir.flow_points[0])/x_num) N = round(len(self.envir.flow_points[1])/y_num) # get worse case max velocity vector for scaling max_u = self.envir.flow[0].max(); max_v = self.envir.flow[1].max() max_mag = np.linalg.norm(np.array([max_u,max_v])) x_pts = self.envir.flow_points[0][::M] y_pts = self.envir.flow_points[1][::N] fld = ax.quiver(x_pts, y_pts, np.zeros((len(y_pts),len(x_pts))), np.zeros((len(y_pts),len(x_pts))), scale=max_mag*5, alpha=0.2) # scatter plot scat = ax.scatter([], [], label='organism', c='darkgreen', s=3) # textual info time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes, fontsize=12) perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_swrm_vel = \ self._calc_basic_stats(DIM3=False, t_indx=n0) axStats = plt.axes([0.77, 0.77, 0.25, 0.2], frameon=False) axStats.set_axis_off() stats_text = axStats.text(0,1, '{:.1f}% remain\n'.format(perc_left)+ '\n ------ Info ------\n'+ r'Fluid $v_{max}$'+': {:.1g} {}/s\n'.format(max_spd, self.envir.units)+ r'Fluid $\overline{v}$'+': {:.1g} {}/s\n'.format(avg_spd, self.envir.units)+ r'Agent $\overline{v}$'+': {:.1g} {}/s\n'.format(np.linalg.norm(avg_swrm_vel), self.envir.units), fontsize=10, transform=axStats.transAxes, verticalalignment='top') x_text = axHistx.text(0.01, 0.98, r'Fluid $\overline{v}_x$'+': {:.2g} \n'.format(avg_spd_x)+ r'Agent $\overline{v}_x$'+': {:.2g}'.format(avg_swrm_vel[0]), transform=axHistx.transAxes, verticalalignment='top', fontsize=10) y_text = axHisty.text(0.02, 0.99, r'Fluid $\overline{v}_y$'+': {:.2g} \n'.format(avg_spd_y)+ r'Agent $\overline{v}_y$'+': {:.2g}'.format(avg_swrm_vel[1]), transform=axHisty.transAxes, verticalalignment='top', fontsize=10) if dist == 'hist': # histograms data_x = self.pos_history[n0][:,0].compressed() data_y = self.pos_history[n0][:,1].compressed() bins_x = np.linspace(0, self.envir.L[0], 26) bins_y = np.linspace(0, self.envir.L[1], 26) n_x, bins_x, patches_x = axHistx.hist(data_x, bins=bins_x) n_y, bins_y, patches_y = axHisty.hist(data_y, bins=bins_y, orientation='horizontal') else: # Gaussian Kernel Density Estimation if dist == 'cov': fac_x = self.shared_props['cov'][0,0] fac_y = self.shared_props['cov'][1,1] else: try: fac_x = float(dist) fac_y = fac_x except: # estimate covariance from Scotts Factor. HOWEVER: this # estimation breaks if IC is point source. fac_x = None fac_y = None xmesh = np.linspace(0, self.envir.L[0], 1000) ymesh = np.linspace(0, self.envir.L[1], 1000) # deal with point sources pos_x = self.pos_history[n0][:,0].compressed() pos_y = self.pos_history[n0][:,1].compressed() try: if len(pos_x) > 1: x_density = stats.gaussian_kde(pos_x, fac_x) x_density = x_density(xmesh) elif len(pos_x) == 1: raise np.linalg.LinAlgError else: x_density = np.zeros_like(xmesh) except np.linalg.LinAlgError: idx = (np.abs(xmesh - pos_x[0])).argmin() x_density = np.zeros_like(xmesh); x_density[idx] = 1 try: if len(pos_y) > 1: y_density = stats.gaussian_kde(pos_y, fac_y) y_density = y_density(ymesh) elif len(pos_y) == 1: raise np.linalg.LinAlgError else: y_density = np.zeros_like(ymesh) except np.linalg.LinAlgError: idy = (np.abs(ymesh - pos_y[0])).argmin() y_density = np.zeros_like(ymesh); y_density[idy] = 1 xdens_plt, = axHistx.plot(xmesh, x_density) ydens_plt, = axHisty.plot(y_density, ymesh) axHistx.get_yaxis().set_ticks([]) axHisty.get_xaxis().set_ticks([]) if np.max(xdens_plt.get_ydata()) != 0: axHistx.set_ylim(bottom=0, top=np.max(xdens_plt.get_ydata())) else: axHistx.set_ylim(bottom=0) if np.max(ydens_plt.get_xdata()) != 0: axHisty.set_xlim(left=0, right=np.max(ydens_plt.get_xdata())) else: axHisty.set_xlim(left=0) else: ### 3D setup ### if figsize is None: fig = plt.figure(figsize=(10,5)) else: fig = plt.figure(figsize=figsize) ax, axHistx, axHisty, axHistz = self.envir._plot_setup(fig) if azim is not None or elev is not None: ax.view_init(elev, azim) if downsamp is None: scat = ax.scatter(self.pos_history[n0][:,0], self.pos_history[n0][:,1], self.pos_history[n0][:,2], label='organism', animated=True) else: scat = ax.scatter(self.pos_history[n0][downsamp,0], self.pos_history[n0][downsamp,1], self.pos_history[n0][downsamp,2], label='organism', animated=True) # textual info time_text = ax.text2D(0.02, 1, 'time = {:.2f}'.format( self.envir.time_history[n0]), transform=ax.transAxes, animated=True, verticalalignment='top', fontsize=12) perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_spd_z, avg_swrm_vel = \ self._calc_basic_stats(DIM3=True, t_indx=n0) flow_text = ax.text2D(0.75, 0.9, r'Fluid $v_{max}$'+': {:.2g} {}/s\n'.format( max_spd, self.envir.units)+ r'Fluid $v_{avg}$'+': {:.2g} {}/s\n'.format( avg_spd, self.envir.units)+ r'Agent $v_{avg}$'+': {:.2g} {}/s'.format( np.linalg.norm(avg_swrm_vel), self.envir.units), transform=ax.transAxes, animated=True, horizontalalignment='left', fontsize=10) perc_text = ax.text2D(0.02, 0, '{:.1f}% remain\n'.format(perc_left), transform=fig.transFigure, animated=True, fontsize=10) x_text = axHistx.text(0.02, 0.98, r'Fluid $\overline{v}_x$'+': {:.2g} {}/s\n'.format( avg_spd_x, self.envir.units)+ r'Agent $\overline{v}_x$'+': {:.2g} {}/s'.format( avg_swrm_vel[0], self.envir.units), transform=axHistx.transAxes, animated=True, verticalalignment='top', fontsize=10) y_text = axHisty.text(0.02, 0.98, r'Fluid $\overline{v}_y$'+': {:.2g} {}/s\n'.format( avg_spd_y, self.envir.units)+ r'Agent $\overline{v}_y$'+': {:.2g} {}/s'.format( avg_swrm_vel[1], self.envir.units), transform=axHisty.transAxes, animated=True, verticalalignment='top', fontsize=10) z_text = axHistz.text(0.02, 0.98, r'Fluid $\overline{v}_z$'+': {:.2g} {}/s\n'.format( avg_spd_z, self.envir.units)+ r'Agent $\overline{v}_z$'+': {:.2g} {}/s'.format( avg_swrm_vel[2], self.envir.units), transform=axHistz.transAxes, animated=True, verticalalignment='top', fontsize=10) if dist == 'hist': # histograms data_x = self.pos_history[n0][:,0].compressed() data_y = self.pos_history[n0][:,1].compressed() data_z = self.pos_history[n0][:,2].compressed() bins_x = np.linspace(0, self.envir.L[0], 26) bins_y = np.linspace(0, self.envir.L[1], 26) bins_z = np.linspace(0, self.envir.L[2], 26) n_x, bins_x, patches_x = axHistx.hist(data_x, bins=bins_x, alpha=0.8) n_y, bins_y, patches_y = axHisty.hist(data_y, bins=bins_y, alpha=0.8) n_z, bins_z, patches_z = axHistz.hist(data_z, bins=bins_z, alpha=0.8) else: # Gaussian Kernel Density Estimation if dist == 'cov': fac_x = self.shared_props['cov'][0,0] fac_y = self.shared_props['cov'][1,1] fac_z = self.shared_props['cov'][2,2] else: try: # see if a float was passed fac_x = float(dist) fac_y = fac_x fac_z = fac_x except: # estimate covariance from Scotts Factor. HOWEVER: this # estimation breaks if IC is point source. fac_x = None fac_y = None fac_z = None xmesh = np.linspace(0, self.envir.L[0], 1000) ymesh = np.linspace(0, self.envir.L[1], 1000) zmesh = np.linspace(0, self.envir.L[2], 1000) # deal with point sources pos_x = self.pos_history[n0][:,0].compressed() pos_y = self.pos_history[n0][:,1].compressed() pos_z = self.pos_history[n0][:,2].compressed() try: if len(pos_x) > 1: x_density = stats.gaussian_kde(pos_x, fac_x) x_density = x_density(xmesh) elif len(pos_x) == 1: raise np.linalg.LinAlgError else: x_density = np.zeros_like(xmesh) except np.linalg.LinAlgError: idx = (np.abs(xmesh - pos_x[0])).argmin() x_density = np.zeros_like(xmesh); x_density[idx] = 1 try: if len(pos_y) > 1: y_density = stats.gaussian_kde(pos_y, fac_y) y_density = y_density(ymesh) elif len(pos_y) == 1: raise np.linalg.LinAlgError else: y_density = np.zeros_like(ymesh) except np.linalg.LinAlgError: idy = (np.abs(ymesh - pos_y[0])).argmin() y_density = np.zeros_like(ymesh); y_density[idy] = 1 try: if len(pos_z) > 1: z_density = stats.gaussian_kde(pos_z, fac_z) z_density = z_density(zmesh) elif len(pos_z) == 1: raise np.linalg.LinAlgError else: z_density = np.zeros_like(zmesh) except np.linalg.LinAlgError: idz = (np.abs(zmesh - pos_z[0])).argmin() z_density = np.zeros_like(zmesh); z_density[idz] = 1 xdens_plt, = axHistx.plot(xmesh, x_density) ydens_plt, = axHisty.plot(ymesh, y_density) zdens_plt, = axHistz.plot(zmesh, z_density) axHistx.get_yaxis().set_ticks([]) axHisty.get_yaxis().set_ticks([]) axHistz.get_yaxis().set_ticks([]) if np.max(xdens_plt.get_ydata()) != 0: axHistx.set_ylim(bottom=0, top=np.max(xdens_plt.get_ydata())) else: axHistx.set_ylim(bottom=0) if np.max(ydens_plt.get_ydata()) != 0: axHisty.set_ylim(bottom=0, top=np.max(ydens_plt.get_ydata())) else: axHisty.set_ylim(bottom=0) if np.max(zdens_plt.get_ydata()) != 0: axHistz.set_ylim(bottom=0, top=np.max(zdens_plt.get_ydata())) else: axHistz.set_ylim(bottom=0) # animation function. Called sequentially def animate(n): if n < len(self.pos_history): time_text.set_text('time = {:.2f}'.format(self.envir.time_history[n])) if not DIM3: # 2D perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_swrm_vel = \ self._calc_basic_stats(DIM3=False, t_indx=n) stats_text.set_text('{:.1f}% remain\n'.format(perc_left)+ '\n ------ Info ------\n'+ r'Fluid $v_{max}$'+': {:.1g} {}/s\n'.format(max_spd, self.envir.units)+ r'Fluid $\overline{v}$'+': {:.1g} {}/s\n'.format(avg_spd, self.envir.units)+ r'Agent $\overline{v}$'+': {:.1g} {}/s\n'.format(np.linalg.norm(avg_swrm_vel), self.envir.units)) x_text.set_text(r'Fluid $\overline{v}_x$'+': {:.2g} \n'.format(avg_spd_x)+ r'Agent $\overline{v}_x$'+': {:.2g}'.format(avg_swrm_vel[0])) y_text.set_text(r'Fluid $\overline{v}_y$'+': {:.2g} \n'.format(avg_spd_y)+ r'Agent $\overline{v}_y$'+': {:.2g}'.format(avg_swrm_vel[1])) if fluid == 'vort' and self.envir.flow is not None: vort = self.envir.get_2D_vorticity(t_indx=n) fld.set_array(vort.T) fld.changed() fld.autoscale() elif fluid == 'quiver' and self.envir.flow is not None: if self.envir.flow_times is not None: flow = self.envir.interpolate_temporal_flow(t_indx=n) fld.set_UVC(flow[0][::M,::N].T, flow[1][::M,::N].T) else: fld.set_UVC(self.envir.flow[0][::M,::N].T, self.envir.flow[1][::M,::N].T) if downsamp is None: scat.set_offsets(self.pos_history[n]) else: scat.set_offsets(self.pos_history[n][downsamp,:]) if dist == 'hist': n_x, _ = np.histogram(self.pos_history[n][:,0].compressed(), bins_x) n_y, _ = np.histogram(self.pos_history[n][:,1].compressed(), bins_y) for rect, h in zip(patches_x, n_x): rect.set_height(h) for rect, h in zip(patches_y, n_y): rect.set_width(h) if fluid == 'vort' and self.envir.flow is not None: return [fld, scat, time_text, stats_text, x_text, y_text] + list(patches_x) + list(patches_y) else: return [scat, time_text, stats_text, x_text, y_text] + list(patches_x) + list(patches_y) else: pos_x = self.pos_history[n][:,0].compressed() pos_y = self.pos_history[n][:,1].compressed() try: if len(pos_x) > 1: x_density = stats.gaussian_kde(pos_x, fac_x) x_density = x_density(xmesh) elif len(pos_x) == 1: raise np.linalg.LinAlgError else: x_density = np.zeros_like(xmesh) except np.linalg.LinAlgError: idx = (np.abs(xmesh - pos_x[0])).argmin() x_density = np.zeros_like(xmesh); x_density[idx] = 1 try: if len(pos_y) > 1: y_density = stats.gaussian_kde(pos_y, fac_y) y_density = y_density(ymesh) elif len(pos_y) == 1: raise np.linalg.LinAlgError else: y_density = np.zeros_like(ymesh) except np.linalg.LinAlgError: idy = (np.abs(ymesh - pos_y[0])).argmin() y_density = np.zeros_like(ymesh); y_density[idy] = 1 xdens_plt.set_ydata(x_density) ydens_plt.set_xdata(y_density) if np.max(xdens_plt.get_ydata()) != 0: axHistx.set_ylim(bottom=0, top=np.max(xdens_plt.get_ydata())) else: axHistx.set_ylim(bottom=0) if np.max(ydens_plt.get_xdata()) != 0: axHisty.set_xlim(left=0, right=np.max(ydens_plt.get_xdata())) else: axHisty.set_xlim(left=0) if fluid == 'vort' and self.envir.flow is not None: return [fld, scat, time_text, stats_text, x_text, y_text, xdens_plt, ydens_plt] else: return [scat, time_text, stats_text, x_text, y_text, xdens_plt, ydens_plt] else: # 3D perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_spd_z, avg_swrm_vel = \ self._calc_basic_stats(DIM3=True, t_indx=n) # print(n) # print(self.pos_history[n].all() is ma.masked) flow_text.set_text(r'Fluid $v_{max}$'+': {:.2g} {}/s\n'.format( max_spd, self.envir.units)+ r'Fluid $v_{avg}$'+': {:.2g} {}/s\n'.format( avg_spd, self.envir.units)+ r'Agent $v_{avg}$'+': {:.2g} {}/s'.format( np.linalg.norm(avg_swrm_vel), self.envir.units)) perc_text.set_text('{:.1f}% remain\n'.format(perc_left)) x_text.set_text(r'Fluid $\overline{v}_x$'+': {:.2g} {}/s\n'.format( avg_spd_x, self.envir.units)+ r'Agent $\overline{v}_x$'+': {:.2g} {}/s'.format( avg_swrm_vel[0], self.envir.units)) y_text.set_text(r'Fluid $\overline{v}_y$'+': {:.2g} {}/s\n'.format( avg_spd_y, self.envir.units)+ r'Agent $\overline{v}_y$'+': {:.2g} {}/s'.format( avg_swrm_vel[1], self.envir.units)) z_text.set_text(r'Fluid $\overline{v}_z$'+': {:.2g} {}/s\n'.format( avg_spd_z, self.envir.units)+ r'Agent $\overline{v}_z$'+': {:.2g} {}/s'.format( avg_swrm_vel[2], self.envir.units)) if downsamp is None: scat._offsets3d = (np.ma.ravel(self.pos_history[n][:,0].compressed()), np.ma.ravel(self.pos_history[n][:,1].compressed()), np.ma.ravel(self.pos_history[n][:,2].compressed())) else: scat._offsets3d = (np.ma.ravel(self.pos_history[n][downsamp,0].compressed()), np.ma.ravel(self.pos_history[n][downsamp,1].compressed()), np.ma.ravel(self.pos_history[n][downsamp,2].compressed())) if dist == 'hist': n_x, _ = np.histogram(self.pos_history[n][:,0].compressed(), bins_x) n_y, _ = np.histogram(self.pos_history[n][:,1].compressed(), bins_y) n_z, _ = np.histogram(self.pos_history[n][:,2].compressed(), bins_z) for rect, h in zip(patches_x, n_x): rect.set_height(h) for rect, h in zip(patches_y, n_y): rect.set_height(h) for rect, h in zip(patches_z, n_z): rect.set_height(h) fig.canvas.draw() return [scat, time_text, flow_text, perc_text, x_text, y_text, z_text] + list(patches_x) + list(patches_y) + list(patches_z) else: pos_x = self.pos_history[n][:,0].compressed() pos_y = self.pos_history[n][:,1].compressed() pos_z = self.pos_history[n][:,2].compressed() try: if len(pos_x) > 1: x_density = stats.gaussian_kde(pos_x, fac_x) x_density = x_density(xmesh) elif len(pos_x) == 1: raise np.linalg.LinAlgError else: x_density = np.zeros_like(xmesh) except np.linalg.LinAlgError: idx = (np.abs(xmesh - pos_x[0])).argmin() x_density = np.zeros_like(xmesh); x_density[idx] = 1 try: if len(pos_y) > 1: y_density = stats.gaussian_kde(pos_y, fac_y) y_density = y_density(ymesh) elif len(pos_y) == 1: raise np.linalg.LinAlgError else: y_density = np.zeros_like(ymesh) except np.linalg.LinAlgError: idy = (np.abs(ymesh - pos_y[0])).argmin() y_density = np.zeros_like(ymesh); y_density[idy] = 1 try: if len(pos_z) > 1: z_density = stats.gaussian_kde(pos_z, fac_z) z_density = z_density(zmesh) elif len(pos_z) == 1: raise np.linalg.LinAlgError else: z_density = np.zeros_like(zmesh) except np.linalg.LinAlgError: idz = (np.abs(zmesh - pos_z[0])).argmin() z_density = np.zeros_like(zmesh); z_density[idz] = 1 xdens_plt.set_ydata(x_density) ydens_plt.set_ydata(y_density) zdens_plt.set_ydata(z_density) if np.max(xdens_plt.get_ydata()) != 0: axHistx.set_ylim(bottom=0, top=np.max(xdens_plt.get_ydata())) else: axHistx.set_ylim(bottom=0) if np.max(ydens_plt.get_ydata()) != 0: axHisty.set_ylim(bottom=0, top=np.max(ydens_plt.get_ydata())) else: axHisty.set_ylim(bottom=0) if np.max(zdens_plt.get_ydata()) != 0: axHistz.set_ylim(bottom=0, top=np.max(zdens_plt.get_ydata())) else: axHistz.set_ylim(bottom=0) fig.canvas.draw() return [scat, time_text, flow_text, perc_text, x_text, y_text, z_text, xdens_plt, ydens_plt, zdens_plt] else: time_text.set_text('time = {:.2f}'.format(self.envir.time)) if not DIM3: # 2D end perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_swrm_vel = \ self._calc_basic_stats(DIM3=False, t_indx=None) stats_text.set_text('{:.1f}% remain\n'.format(perc_left)+ '\n ------ Info ------\n'+ r'Fluid $v_{max}$'+': {:.1g} {}/s\n'.format(max_spd, self.envir.units)+ r'Fluid $\overline{v}$'+': {:.1g} {}/s\n'.format(avg_spd, self.envir.units)+ r'Agent $\overline{v}$'+': {:.1g} {}/s\n'.format(np.linalg.norm(avg_swrm_vel), self.envir.units)) x_text.set_text(r'Fluid $\overline{v}_x$'+': {:.2g} \n'.format(avg_spd_x)+ r'Agent $\overline{v}_x$'+': {:.2g}'.format(avg_swrm_vel[0])) y_text.set_text(r'Fluid $\overline{v}_y$'+': {:.2g} \n'.format(avg_spd_y)+ r'Agent $\overline{v}_y$'+': {:.2g}'.format(avg_swrm_vel[1])) if fluid == 'vort' and self.envir.flow is not None: vort = self.envir.get_2D_vorticity() fld.set_array(vort.T) fld.changed() fld.autoscale() elif fluid == 'quiver' and self.envir.flow is not None: if self.envir.flow_times is not None: flow = self.envir.interpolate_temporal_flow() fld.set_UVC(flow[0][::M,::N].T, flow[1][::M,::N].T) else: fld.set_UVC(self.envir.flow[0][::M,::N].T, self.envir.flow[1][::M,::N].T) if downsamp is None: scat.set_offsets(self.positions) else: scat.set_offsets(self.positions[downsamp,:]) if dist == 'hist': n_x, _ = np.histogram(self.positions[:,0].compressed(), bins_x) n_y, _ = np.histogram(self.positions[:,1].compressed(), bins_y) for rect, h in zip(patches_x, n_x): rect.set_height(h) for rect, h in zip(patches_y, n_y): rect.set_width(h) if fluid == 'vort' and self.envir.flow is not None: return [fld, scat, time_text, stats_text, x_text, y_text] + list(patches_x) + list(patches_y) else: return [scat, time_text, stats_text, x_text, y_text] + list(patches_x) + list(patches_y) else: pos_x = self.positions[:,0].compressed() pos_y = self.positions[:,1].compressed() try: if len(pos_x) > 1: x_density = stats.gaussian_kde(pos_x, fac_x) x_density = x_density(xmesh) elif len(pos_x) == 1: raise np.linalg.LinAlgError else: x_density = np.zeros_like(xmesh) except np.linalg.LinAlgError: idx = (np.abs(xmesh - pos_x[0])).argmin() x_density = np.zeros_like(xmesh); x_density[idx] = 1 try: if len(pos_y) > 1: y_density = stats.gaussian_kde(pos_y, fac_y) y_density = y_density(ymesh) elif len(pos_y) == 1: raise np.linalg.LinAlgError else: y_density = np.zeros_like(ymesh) except np.linalg.LinAlgError: idy = (np.abs(ymesh - pos_y[0])).argmin() y_density = np.zeros_like(ymesh); y_density[idy] = 1 xdens_plt.set_ydata(x_density) ydens_plt.set_xdata(y_density) if np.max(xdens_plt.get_ydata()) != 0: axHistx.set_ylim(bottom=0, top=np.max(xdens_plt.get_ydata())) else: axHistx.set_ylim(bottom=0) if np.max(ydens_plt.get_xdata()) != 0: axHisty.set_xlim(left=0, right=np.max(ydens_plt.get_xdata())) else: axHisty.set_xlim(left=0) if fluid == 'vort' and self.envir.flow is not None: return [fld, scat, time_text, stats_text, x_text, y_text, xdens_plt, ydens_plt] else: return [scat, time_text, stats_text, x_text, y_text, xdens_plt, ydens_plt] else: # 3D end perc_left, avg_spd, max_spd, avg_spd_x, avg_spd_y, avg_spd_z, avg_swrm_vel = \ self._calc_basic_stats(DIM3=True) flow_text.set_text(r'Fluid $v_{max}$'+': {:.2g} {}/s\n'.format( max_spd, self.envir.units)+ r'Fluid $v_{avg}$'+': {:.2g} {}/s\n'.format( avg_spd, self.envir.units)+ r'Agent $v_{avg}$'+': {:.2g} {}/s'.format( np.linalg.norm(avg_swrm_vel), self.envir.units)) perc_text.set_text('{:.1f}% remain\n'.format(perc_left)) x_text.set_text(r'Fluid $\overline{v}_x$'+': {:.2g} {}/s\n'.format( avg_spd_x, self.envir.units)+ r'Agent $\overline{v}_x$'+': {:.2g} {}/s'.format( avg_swrm_vel[0], self.envir.units)) y_text.set_text(r'Fluid $\overline{v}_y$'+': {:.2g} {}/s\n'.format( avg_spd_y, self.envir.units)+ r'Agent $\overline{v}_y$'+': {:.2g} {}/s'.format( avg_swrm_vel[1], self.envir.units)) z_text.set_text(r'Fluid $\overline{v}_z$'+': {:.2g} {}/s\n'.format( avg_spd_z, self.envir.units)+ r'Agent $\overline{v}_z$'+': {:.2g} {}/s'.format( avg_swrm_vel[2], self.envir.units)) if downsamp is None: scat._offsets3d = (np.ma.ravel(self.positions[:,0].compressed()), np.ma.ravel(self.positions[:,1].compressed()), np.ma.ravel(self.positions[:,2].compressed())) else: scat._offsets3d = (np.ma.ravel(self.positions[downsamp,0].compressed()), np.ma.ravel(self.positions[downsamp,1].compressed()), np.ma.ravel(self.positions[downsamp,2].compressed())) if dist == 'hist': n_x, _ = np.histogram(self.positions[:,0].compressed(), bins_x) n_y, _ = np.histogram(self.positions[:,1].compressed(), bins_y) n_z, _ = np.histogram(self.positions[:,2].compressed(), bins_z) for rect, h in zip(patches_x, n_x): rect.set_height(h) for rect, h in zip(patches_y, n_y): rect.set_height(h) for rect, h in zip(patches_z, n_z): rect.set_height(h) fig.canvas.draw() return [scat, time_text, flow_text, perc_text, x_text, y_text, z_text] + list(patches_x) + list(patches_y) + list(patches_z) else: pos_x = self.positions[:,0].compressed() pos_y = self.positions[:,1].compressed() pos_z = self.positions[:,2].compressed() try: if len(pos_x) > 1: x_density = stats.gaussian_kde(pos_x, fac_x) x_density = x_density(xmesh) elif len(pos_x) == 1: raise np.linalg.LinAlgError else: x_density = np.zeros_like(xmesh) except np.linalg.LinAlgError: idx = (np.abs(xmesh - pos_x[0])).argmin() x_density = np.zeros_like(xmesh); x_density[idx] = 1 try: if len(pos_y) > 1: y_density = stats.gaussian_kde(pos_y, fac_y) y_density = y_density(ymesh) elif len(pos_y) == 1: raise np.linalg.LinAlgError else: y_density = np.zeros_like(ymesh) except np.linalg.LinAlgError: idy = (np.abs(ymesh - pos_y[0])).argmin() y_density = np.zeros_like(ymesh); y_density[idy] = 1 try: if len(pos_z) > 1: z_density = stats.gaussian_kde(pos_z, fac_z) z_density = z_density(zmesh) elif len(pos_z) == 1: raise np.linalg.LinAlgError else: z_density = np.zeros_like(zmesh) except np.linalg.LinAlgError: idz = (np.abs(zmesh - pos_z[0])).argmin() z_density = np.zeros_like(zmesh); z_density[idz] = 1 xdens_plt.set_ydata(x_density) ydens_plt.set_ydata(y_density) zdens_plt.set_ydata(z_density) if np.max(xdens_plt.get_ydata()) != 0: axHistx.set_ylim(bottom=0, top=np.max(xdens_plt.get_ydata())) else: axHistx.set_ylim(bottom=0) if np.max(ydens_plt.get_ydata()) != 0: axHisty.set_ylim(bottom=0, top=np.max(ydens_plt.get_ydata())) else: axHisty.set_ylim(bottom=0) if np.max(zdens_plt.get_ydata()) != 0: axHistz.set_ylim(bottom=0, top=np.max(zdens_plt.get_ydata())) else: axHistz.set_ylim(bottom=0) fig.canvas.draw() return [scat, time_text, flow_text, perc_text, x_text, y_text, z_text, xdens_plt, ydens_plt, zdens_plt] # infer animation rate from dt between current and last position dt = self.envir.time - self.envir.time_history[-1] if frames is None: frames = range(len(self.pos_history)+1) anim = animation.FuncAnimation(fig, animate, frames=frames, interval=dt*100, repeat=False, blit=True, save_count=len(frames)) if movie_filename is not None: try: if writer_kwargs is None: writer = animation.FFMpegWriter(fps=fps, metadata=dict(artist='Christopher Strickland'))#, bitrate=1800) else: writer = animation.FFMpegWriter(**writer_kwargs) if save_kwargs is None: anim.save(movie_filename, writer=writer, dpi=150) else: anim.save(movie_filename, writer=writer, **save_kwargs) plt.close() print('Video saved to {}.'.format(movie_filename)) except: print('Failed to save animation.') print('Check that you have ffmpeg or mencoder installed; these') print("aren't Python packages, but stand-alone applications.") print("An H.264 encoder is needed on the system's path in order") print('to save to that video format.') raise else: plt.show()