<s> from __future__ import print_function <EOL> import os <EOL> import math <EOL> import warnings <EOL> import numpy as np <EOL> from matplotlib . collections import LineCollection <EOL> from matplotlib . transforms import Bbox <EOL> from matplotlib . patches import Polygon <EOL> from . geometry . path import Path , get_endpoints <EOL> from . import extract_pv_slice <EOL> def distance ( x1 , y1 , x2 , y2 , x3 , y3 ) : <EOL> """<STR_LIT>""" <EOL> px = x2 - x1 <EOL> py = y2 - y1 <EOL> something = px * px + py * py <EOL> u = ( ( x3 - x1 ) * px + ( y3 - y1 ) * py ) / float ( something ) <EOL> x = x1 + u * px <EOL> y = y1 + u * py <EOL> dx = x - x3 <EOL> dy = y - y3 <EOL> dist = math . sqrt ( dx * dx + dy * dy ) <EOL> return dist <EOL> class MovableSliceBox ( object ) : <EOL> def __init__ ( self , box , callback ) : <EOL> self . box = box <EOL> self . press = None <EOL> self . background = None <EOL> self . point_counter = <NUM_LIT:0> <EOL> self . callback = callback <EOL> self . mode = <NUM_LIT:0> <EOL> self . show_poly = False <EOL> self . cidpress = self . box . figure . canvas . mpl_connect ( '<STR_LIT>' , self . draw_slicer ) <EOL> def connect ( self ) : <EOL> self . cidpress = self . box . figure . canvas . mpl_connect ( '<STR_LIT>' , self . key_press ) <EOL> self . cidpress = self . box . figure . canvas . mpl_connect ( '<STR_LIT>' , self . on_press ) <EOL> self . cidmotion = self . box . figure . canvas . mpl_connect ( '<STR_LIT>' , self . on_motion ) <EOL> def draw_slicer ( self , event ) : <EOL> axes = self . box . axes <EOL> canvas = self . box . figure . canvas <EOL> self . box . axes . draw_artist ( self . box ) <EOL> if self . show_poly : <EOL> path = Path ( zip ( self . box . x , self . box . y ) ) <EOL> path . width = self . box . width <EOL> patches = path . to_patches ( <NUM_LIT:1> , ec = '<STR_LIT>' , fc = '<STR_LIT:none>' , <EOL> transform = self . box . axes . transData , <EOL> clip_on = True , clip_box = self . box . axes . bbox ) <EOL> for patch in patches : <EOL> self . box . axes . draw_artist ( patch ) <EOL> def on_press ( self , event ) : <EOL> if self . box . figure . canvas . toolbar . mode != '<STR_LIT>' : <EOL> return <EOL> if event . inaxes != self . box . axes : <EOL> return <EOL> if self . mode == <NUM_LIT:1> : <EOL> self . callback ( self . box ) <EOL> self . mode += <NUM_LIT:1> <EOL> return <EOL> if self . mode == <NUM_LIT:2> : <EOL> self . box . x = [ ] <EOL> self . box . y = [ ] <EOL> self . mode = <NUM_LIT:0> <EOL> self . point_counter = <NUM_LIT:0> <EOL> self . press = event . xdata , event . ydata <EOL> self . point_counter += <NUM_LIT:1> <EOL> axes = self . box . axes <EOL> canvas = self . box . figure . canvas <EOL> if self . point_counter == <NUM_LIT:1> : <EOL> self . box . x . append ( event . xdata ) <EOL> self . box . x . append ( event . xdata ) <EOL> self . box . y . append ( event . ydata ) <EOL> self . box . y . append ( event . ydata ) <EOL> self . box . width = <NUM_LIT:0.> <EOL> self . box . set_animated ( True ) <EOL> canvas . draw ( ) <EOL> self . background = canvas . copy_from_bbox ( self . box . axes . bbox ) <EOL> elif self . mode == <NUM_LIT:0> : <EOL> self . box . x . append ( event . xdata ) <EOL> self . box . y . append ( event . ydata ) <EOL> self . box . _update_segments ( ) <EOL> axes . draw_artist ( self . box ) <EOL> canvas . blit ( axes . bbox ) <EOL> def key_press ( self , event ) : <EOL> if self . box . figure . canvas . toolbar . mode != '<STR_LIT>' : <EOL> return <EOL> if event . key == '<STR_LIT>' and self . mode == <NUM_LIT:0> : <EOL> self . mode += <NUM_LIT:1> <EOL> self . box . x = self . box . x [ : - <NUM_LIT:1> ] <EOL> self . box . y = self . box . y [ : - <NUM_LIT:1> ] <EOL> if event . key == '<STR_LIT:y>' and self . mode == <NUM_LIT:2> : <EOL> self . show_poly = not self . show_poly <EOL> self . draw_slicer ( event ) <EOL> self . box . figure . canvas . draw ( ) <EOL> def on_motion ( self , event ) : <EOL> if self . box . figure . canvas . toolbar . mode != '<STR_LIT>' : <EOL> return <EOL> if self . point_counter == <NUM_LIT:0> : <EOL> return <EOL> if self . mode == <NUM_LIT:2> : <EOL> return <EOL> canvas = self . box . figure . canvas <EOL> axes = self . box . axes <EOL> canvas . restore_region ( self . background ) <EOL> if event . inaxes != self . box . axes : <EOL> return <EOL> if self . mode == <NUM_LIT:0> : <EOL> self . box . x [ - <NUM_LIT:1> ] = event . xdata <EOL> self . box . y [ - <NUM_LIT:1> ] = event . ydata <EOL> elif self . mode == <NUM_LIT:1> : <EOL> self . box . width = distance ( self . box . x [ - <NUM_LIT:2> ] , self . box . y [ - <NUM_LIT:2> ] , self . box . x [ - <NUM_LIT:1> ] , self . box . y [ - <NUM_LIT:1> ] , event . xdata , event . ydata ) * <NUM_LIT:2> <EOL> self . box . _update_segments ( ) <EOL> axes . draw_artist ( self . box ) <EOL> canvas . blit ( axes . bbox ) <EOL> def disconnect ( self ) : <EOL> self . box . figure . canvas . mpl_disconnect ( self . cidpress ) <EOL> self . box . figure . canvas . mpl_disconnect ( self . cidmotion ) <EOL> class SliceCurve ( LineCollection ) : <EOL> def __init__ ( self , x = [ ] , y = [ ] , width = None , ** kwargs ) : <EOL> super ( SliceCurve , self ) . __init__ ( [ ] , ** kwargs ) <EOL> self . x = x <EOL> self . y = y <EOL> self . width = width <EOL> self . _update_segments ( ) <EOL> def _update_segments ( self ) : <EOL> if not self . x : <EOL> return <EOL> x1 , y1 , x2 , y2 = get_endpoints ( self . x , self . y , self . width ) <EOL> line = zip ( self . x , self . y ) <EOL> rect = zip ( np . hstack ( [ x1 , x2 [ : : - <NUM_LIT:1> ] , x1 [ <NUM_LIT:0> ] ] ) , <EOL> np . hstack ( [ y1 , y2 [ : : - <NUM_LIT:1> ] , y1 [ <NUM_LIT:0> ] ] ) ) <EOL> self . set_segments ( ( list ( line ) , list ( rect ) ) ) <EOL> self . set_linestyles ( ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> self . set_linewidths ( ( <NUM_LIT:2> , <NUM_LIT:1> ) ) <EOL> class PVSlicer ( object ) : <EOL> def __init__ ( self , filename , backend = "<STR_LIT>" , clim = None ) : <EOL> self . filename = filename <EOL> try : <EOL> from spectral_cube import SpectralCube <EOL> cube = SpectralCube . read ( filename , format = '<STR_LIT>' ) <EOL> self . array = cube . _data <EOL> except : <EOL> warnings . warn ( "<STR_LIT>" ) <EOL> from astropy . io import fits <EOL> self . array = fits . getdata ( filename ) <EOL> if self . array . ndim != <NUM_LIT:3> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> self . backend = backend <EOL> import matplotlib as mpl <EOL> mpl . use ( self . backend ) <EOL> import matplotlib . pyplot as plt <EOL> self . fig = plt . figure ( figsize = ( <NUM_LIT> , <NUM_LIT:8> ) ) <EOL> self . ax1 = self . fig . add_axes ( [ <NUM_LIT:0.1> , <NUM_LIT:0.1> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> if clim is None : <EOL> warnings . warn ( "<STR_LIT>" ) <EOL> n1 = max ( self . array . shape [ <NUM_LIT:0> ] / <NUM_LIT:10> , <NUM_LIT:1> ) <EOL> n2 = max ( self . array . shape [ <NUM_LIT:1> ] / <NUM_LIT:10> , <NUM_LIT:1> ) <EOL> n3 = max ( self . array . shape [ <NUM_LIT:2> ] / <NUM_LIT:10> , <NUM_LIT:1> ) <EOL> sub_array = self . array [ : : n1 , : : n2 , : : n3 ] <EOL> cmin = np . min ( sub_array [ ~ np . isnan ( sub_array ) & ~ np . isinf ( sub_array ) ] ) <EOL> cmax = np . max ( sub_array [ ~ np . isnan ( sub_array ) & ~ np . isinf ( sub_array ) ] ) <EOL> crange = cmax - cmin <EOL> self . _clim = ( cmin - crange , cmax + crange ) <EOL> else : <EOL> self . _clim = clim <EOL> self . slice = int ( round ( self . array . shape [ <NUM_LIT:0> ] / <NUM_LIT> ) ) <EOL> from matplotlib . widgets import Slider <EOL> self . slice_slider_ax = self . fig . add_axes ( [ <NUM_LIT:0.1> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> self . slice_slider_ax . set_xticklabels ( "<STR_LIT>" ) <EOL> self . slice_slider_ax . set_yticklabels ( "<STR_LIT>" ) <EOL> self . slice_slider = Slider ( self . slice_slider_ax , "<STR_LIT>" , <NUM_LIT:0> , self . array . shape [ <NUM_LIT:0> ] , valinit = self . slice , valfmt = "<STR_LIT>" ) <EOL> self . slice_slider . on_changed ( self . update_slice ) <EOL> self . slice_slider . drawon = False <EOL> self . image = self . ax1 . imshow ( self . array [ self . slice , : , : ] , origin = '<STR_LIT>' , interpolation = '<STR_LIT>' , vmin = self . _clim [ <NUM_LIT:0> ] , vmax = self . _clim [ <NUM_LIT:1> ] , cmap = plt . cm . gray ) <EOL> self . vmin_slider_ax = self . fig . add_axes ( [ <NUM_LIT:0.1> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> self . vmin_slider_ax . set_xticklabels ( "<STR_LIT>" ) <EOL> self . vmin_slider_ax . set_yticklabels ( "<STR_LIT>" ) <EOL> self . vmin_slider = Slider ( self . vmin_slider_ax , "<STR_LIT>" , self . _clim [ <NUM_LIT:0> ] , self . _clim [ <NUM_LIT:1> ] , valinit = self . _clim [ <NUM_LIT:0> ] ) <EOL> self . vmin_slider . on_changed ( self . update_vmin ) <EOL> self . vmin_slider . drawon = False <EOL> self . vmax_slider_ax = self . fig . add_axes ( [ <NUM_LIT:0.1> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> self . vmax_slider_ax . set_xticklabels ( "<STR_LIT>" ) <EOL> self . vmax_slider_ax . set_yticklabels ( "<STR_LIT>" ) <EOL> self . vmax_slider = Slider ( self . vmax_slider_ax , "<STR_LIT>" , self . _clim [ <NUM_LIT:0> ] , self . _clim [ <NUM_LIT:1> ] , valinit = self . _clim [ <NUM_LIT:1> ] ) <EOL> self . vmax_slider . on_changed ( self . update_vmax ) <EOL> self . vmax_slider . drawon = False <EOL> self . grid1 = None <EOL> self . grid2 = None <EOL> self . grid3 = None <EOL> self . ax2 = self . fig . add_axes ( [ <NUM_LIT> , <NUM_LIT:0.1> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> self . box = SliceCurve ( colors = ( <NUM_LIT> , <NUM_LIT:0.0> , <NUM_LIT:0.0> ) ) <EOL> self . ax1 . add_collection ( self . box ) <EOL> self . movable = MovableSliceBox ( self . box , callback = self . update_pv_slice ) <EOL> self . movable . connect ( ) <EOL> from matplotlib . widgets import Button <EOL> self . save_button_ax = self . fig . add_axes ( [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] ) <EOL> self . save_button = Button ( self . save_button_ax , '<STR_LIT>' ) <EOL> self . save_button . on_clicked ( self . save_fits ) <EOL> self . file_status_text = self . fig . text ( <NUM_LIT> , <NUM_LIT> , "<STR_LIT>" , ha = '<STR_LIT>' , va = '<STR_LIT>' ) <EOL> self . set_file_status ( None ) <EOL> self . set_file_status ( None ) <EOL> self . pv_slice = None <EOL> self . cidpress = self . fig . canvas . mpl_connect ( '<STR_LIT>' , self . click ) <EOL> def set_file_status ( self , status , filename = None ) : <EOL> if status == '<STR_LIT>' : <EOL> self . file_status_text . set_text ( '<STR_LIT>' ) <EOL> self . file_status_text . set_color ( '<STR_LIT>' ) <EOL> elif status == '<STR_LIT>' : <EOL> self . file_status_text . set_text ( '<STR_LIT>' . format ( filename ) ) <EOL> self . file_status_text . set_color ( '<STR_LIT>' ) <EOL> else : <EOL> self . file_status_text . set_text ( '<STR_LIT>' ) <EOL> self . file_status_text . set_color ( '<STR_LIT>' ) <EOL> self . fig . canvas . draw ( ) <EOL> def click ( self , event ) : <EOL> if event . inaxes != self . ax2 : <EOL> return <EOL> self . slice_slider . set_val ( event . ydata ) <EOL> def save_fits ( self , * args , ** kwargs ) : <EOL> self . set_file_status ( '<STR_LIT>' ) <EOL> print ( "<STR_LIT>" , end = '<STR_LIT>' ) <EOL> try : <EOL> plot_name = raw_input ( ) <EOL> except NameError : <EOL> plot_name = input ( ) <EOL> if self . pv_slice is None : <EOL> return <EOL> from astropy . io import fits <EOL> self . pv_slice . writeto ( plot_name , clobber = True ) <EOL> print ( "<STR_LIT>" , plot_name ) <EOL> self . set_file_status ( '<STR_LIT>' , filename = plot_name ) <EOL> def update_pv_slice ( self , box ) : <EOL> path = Path ( zip ( box . x , box . y ) ) <EOL> path . width = box . width <EOL> self . pv_slice = extract_pv_slice ( self . array , path ) <EOL> self . ax2 . cla ( ) <EOL> self . ax2 . imshow ( self . pv_slice . data , origin = '<STR_LIT>' , aspect = '<STR_LIT>' , interpolation = '<STR_LIT>' ) <EOL> self . fig . canvas . draw ( ) <EOL> def show ( self , block = True ) : <EOL> import matplotlib . pyplot as plt <EOL> plt . show ( block = block ) <EOL> def update_slice ( self , pos = None ) : <EOL> if self . array . ndim == <NUM_LIT:2> : <EOL> self . image . set_array ( self . array ) <EOL> else : <EOL> self . slice = int ( round ( pos ) ) <EOL> self . image . set_array ( self . array [ self . slice , : , : ] ) <EOL> self . fig . canvas . draw ( ) <EOL> def update_vmin ( self , vmin ) : <EOL> if vmin > self . _clim [ <NUM_LIT:1> ] : <EOL> self . _clim = ( self . _clim [ <NUM_LIT:1> ] , self . _clim [ <NUM_LIT:1> ] ) <EOL> else : <EOL> self . _clim = ( vmin , self . _clim [ <NUM_LIT:1> ] ) <EOL> self . image . set_clim ( * self . _clim ) <EOL> self . fig . canvas . draw ( ) <EOL> def update_vmax ( self , vmax ) : <EOL> if vmax < self . _clim [ <NUM_LIT:0> ] : <EOL> self . _clim = ( self . _clim [ <NUM_LIT:0> ] , self . _clim [ <NUM_LIT:0> ] ) <EOL> else : <EOL> self . _clim = ( self . _clim [ <NUM_LIT:0> ] , vmax ) <EOL> self . image . set_clim ( * self . _clim ) <EOL> self . fig . canvas . draw ( ) </s>
