Source code for yt.visualization.streamlines

"""
Import the components of the volume rendering extension



"""

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

import numpy as np
from yt.data_objects.construction_data_containers import YTStreamlineBase
from yt.funcs import *
from yt.utilities.parallel_tools.parallel_analysis_interface import \
    ParallelAnalysisInterface, parallel_passthrough
from yt.utilities.amr_kdtree.api import AMRKDTree

def sanitize_length(length, ds):
    # Ensure that lengths passed in with units are returned as code_length
    # magnitudes without units
    if isinstance(length, YTArray):
        return ds.arr(length).in_units('code_length').d
    else:
        return length

[docs]class Streamlines(ParallelAnalysisInterface): r"""A collection of streamlines that flow through the volume The Streamlines object contains a collection of streamlines defined as paths that are parallel to a specified vector field. Parameters ---------- ds : `~yt.data_objects.Dataset` This is the dataset to streamline pos : array_like An array of initial starting positions of the streamlines. xfield: field, optional The x component of the vector field to be streamlined. Default:'velocity_x' yfield: field, optional The y component of the vector field to be streamlined. Default:'velocity_y' zfield: field, optional The z component of the vector field to be streamlined. Default:'velocity_z' volume : `yt.extensions.volume_rendering.AMRKDTree`, optional The volume to be streamlined. Can be specified for finer-grained control, but otherwise will be automatically generated. At this point it must use the AMRKDTree. Default: None dx : float, optional Optionally specify the step size during the integration. Default: minimum dx length : float, optional Optionally specify the length of integration. Default: np.max(self.ds.domain_right_edge-self.ds.domain_left_edge) direction : real, optional Specifies the direction of integration. The magnitude of this value has no effect, only the sign. get_magnitude: bool, optional Specifies whether the Streamlines.magnitudes array should be filled with the magnitude of the vector field at each point in the streamline. This seems to be a ~10% hit to performance. Default: False Examples -------- >>> import yt >>> import numpy as np >>> import matplotlib.pylab as pl >>> >>> from yt.visualization.api import Streamlines >>> from mpl_toolkits.mplot3d import Axes3D >>> >>> # Load the dataset and set some parameters >>> ds = load('IsolatedGalaxy/galaxy0030/galaxy0030') >>> c = np.array([0.5]*3) >>> N = 100 >>> scale = 1.0 >>> pos_dx = np.random.random((N,3))*scale-scale/2. >>> pos = c+pos_dx >>> >>> # Define and construct streamlines >>> streamlines = Streamlines( ds,pos, 'velocity_x', 'velocity_y', 'velocity_z', length=1.0) >>> streamlines.integrate_through_volume() >>> >>> # Make a 3D plot of the streamlines and save it to disk >>> fig=pl.figure() >>> ax = Axes3D(fig) >>> for stream in streamlines.streamlines: >>> stream = stream[np.all(stream != 0.0, axis=1)] >>> ax.plot3D(stream[:,0], stream[:,1], stream[:,2], alpha=0.1) >>> pl.savefig('streamlines.png') """
[docs] def __init__(self, ds, positions, xfield='velocity_x', yfield='velocity_x', zfield='velocity_x', volume=None, dx=None, length=None, direction=1, get_magnitude=False): ParallelAnalysisInterface.__init__(self) self.ds = ds self.start_positions = sanitize_length(positions, ds) self.N = self.start_positions.shape[0] # I need a data object to resolve the field names to field tuples # via _determine_fields() ad = self.ds.all_data() self.xfield = ad._determine_fields(xfield)[0] self.yfield = ad._determine_fields(yfield)[0] self.zfield = ad._determine_fields(zfield)[0] self.get_magnitude=get_magnitude self.direction = np.sign(direction) if volume is None: volume = AMRKDTree(self.ds) volume.set_fields([self.xfield,self.yfield,self.zfield], [False,False,False], False) volume.join_parallel_trees() self.volume = volume if dx is None: dx = self.ds.index.get_smallest_dx() self.dx = sanitize_length(dx, ds) if length is None: length = np.max(self.ds.domain_right_edge-self.ds.domain_left_edge) self.length = sanitize_length(length, ds) self.steps = int(length/dx)+1 # Fix up the dx. self.dx = 1.0*self.length/self.steps self.streamlines = np.zeros((self.N,self.steps,3), dtype='float64') self.magnitudes = None if self.get_magnitude: self.magnitudes = np.zeros((self.N,self.steps), dtype='float64')
[docs] def integrate_through_volume(self): nprocs = self.comm.size my_rank = self.comm.rank self.streamlines[my_rank::nprocs,0,:] = \ self.start_positions[my_rank::nprocs] pbar = get_pbar("Streamlining", self.N) for i,stream in enumerate(self.streamlines[my_rank::nprocs]): thismag = None if self.get_magnitude: thismag = self.magnitudes[i,:] step = self.steps while (step > 1): this_node = self.volume.locate_node(stream[-step,:]) step = self._integrate_through_brick( this_node, stream, step, mag=thismag) pbar.update(i) pbar.finish() self._finalize_parallel(None) self.streamlines = self.ds.arr(self.streamlines, 'code_length') if self.get_magnitude: self.magnitudes = self.ds.arr( self.magnitudes, self.ds.field_info[self.xfield].units)
@parallel_passthrough def _finalize_parallel(self,data): self.streamlines = self.comm.mpi_allreduce(self.streamlines, op='sum') if self.get_magnitude: self.magnitudes = self.comm.mpi_allreduce( self.magnitudes, op='sum') def _integrate_through_brick(self, node, stream, step, periodic=False, mag=None): while (step > 1): self.volume.get_brick_data(node) brick = node.data stream[-step+1] = stream[-step] if mag is None: brick.integrate_streamline( stream[-step+1], self.direction*self.dx, None) else: marr = [mag] brick.integrate_streamline( stream[-step+1], self.direction*self.dx, marr) mag[-step+1] = marr[0] if np.any(stream[-step+1,:] <= self.ds.domain_left_edge) | \ np.any(stream[-step+1,:] >= self.ds.domain_right_edge): return 0 if np.any(stream[-step+1,:] < node.get_left_edge()) | \ np.any(stream[-step+1,:] >= node.get_right_edge()): return step-1 step -= 1 return step
[docs] def clean_streamlines(self): temp = np.empty(self.N, dtype='object') temp2 = np.empty(self.N, dtype='object') for i,stream in enumerate(self.streamlines): mask = np.all(stream != 0.0, axis=1) temp[i] = stream[mask] temp2[i] = self.magnitudes[i,mask] self.streamlines = temp self.magnitudes = temp2
[docs] def path(self, streamline_id): """ Returns an YTSelectionContainer1D object defined by a streamline. Parameters ---------- streamline_id : int This defines which streamline from the Streamlines object that will define the YTSelectionContainer1D object. Returns ------- An YTStreamlineBase YTSelectionContainer1D object Examples -------- >>> from yt.visualization.api import Streamlines >>> streamlines = Streamlines(ds, [0.5]*3) >>> streamlines.integrate_through_volume() >>> stream = streamlines.path(0) >>> matplotlib.pylab.semilogy(stream['t'], stream['Density'], '-x') """ return YTStreamlineBase(self.streamlines[streamline_id], ds=self.ds, length = self.length)