class OrthographicRayBlaster(RayBlaster): center = traittypes.Array().valid(check_dtype("f4"), check_shape(3)) forward = traittypes.Array().valid(check_dtype("f4"), check_shape(3)) up = traittypes.Array().valid(check_dtype("f4"), check_shape(3)) east = traittypes.Array().valid(check_dtype("f4"), check_shape(3)) width = traitlets.CFloat(1.0) height = traitlets.CFloat(1.0) nx = traitlets.CInt(512) ny = traitlets.CInt(512) @traitlets.default("east") def _default_east(self): return np.cross(self.forward, self.up) def __init__(self, *args, **kwargs): super(OrthographicRayBlaster, self).__init__(*args, **kwargs) # here origin is not the center, but the bottom left self._directions = np.zeros((self.nx, self.ny, 3), dtype="f4") self._directions[:] = self.forward[None, None, :] self.directions = self._directions.view().reshape( (self.nx * self.ny, 3)) self._origins = np.zeros((self.nx, self.ny, 3), dtype="f4") offset_x, offset_y = np.mgrid[-self.width / 2:self.width / 2:self.nx * 1j, -self.height / 2:self.height / 2:self.ny * 1j, ] self._origins[:] = (self.center + offset_x[..., None] * self.east + offset_y[..., None] * self.up) self.origins = self._origins.view().reshape((self.nx * self.ny, 3))
class Bus(t.HasTraits): '''Bus Model''' name = t.CUnicode(default_value='Bus1', help='Name of Generator (str)') bus_type = t.Enum( values=['SWING', 'PQ', 'PV'], default_value='PQ', help='Bus type', ) real_power_demand = tt.Array(default_value=[0.0], minlen=1, help='Active power demand (MW)') imag_power_demand = tt.Array(default_value=[0.0], minlen=1, help='Reactive power demand (MVAR)') shunt_conductance = t.CFloat(default_value=0, help='Shunt Conductance (TODO: units)') shunt_susceptance = t.CFloat(default_value=0, help='Shunt Susceptance (TODO: units)') area = t.CUnicode(default_value='0', help='Area the bus is located in') voltage_magnitude = t.CFloat(default_value=1.0, help='Voltage magnitude (p.u.)') voltage_angle = t.CFloat(default_value=0.0, help='Voltage angle (deg)') base_voltage = t.CFloat(default_value=230, help='Base voltage (kV)') zone = t.CUnicode(default_value='0', help='Zone the bus is located in') maximum_voltage = t.CFloat(default_value=1.05, help='Maximum voltage') minimum_voltage = t.CFloat(default_value=0.95, help='Minimum voltage') def __init__(self, **kwargs): v = kwargs.pop('real_power_demand', None) v = self._coerce_value_to_list({'value': v}) kwargs['real_power_demand'] = v v = kwargs.pop('imag_power_demand', None) v = self._coerce_value_to_list({'value': v}) kwargs['imag_power_demand'] = v super(Bus, self).__init__(**kwargs) @t.validate('real_power_demand', 'imag_power_demand') def _coerce_value_to_list(self, proposal): v = proposal['value'] if (v is not None and (isinstance(v, int) or isinstance(v, float))): v = [ v, ] elif v is None: v = [ 0.0, ] return v
class VertexArray(traitlets.HasTraits): name = traitlets.CUnicode("vertex") id = traitlets.CInt(-1) indices = traittypes.Array(None, allow_none=True) index_id = traitlets.CInt(-1) attributes = traitlets.List(trait=traitlets.Instance(VertexAttribute)) each = traitlets.CInt(-1) @traitlets.default("id") def _id_default(self): return GL.glGenVertexArrays(1) @contextmanager def bind(self, program=None): GL.glBindVertexArray(self.id) if self.index_id != -1: GL.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, self.index_id) # We only bind the attributes if we have a program too if program is None: attrs = [] else: attrs = self.attributes with ExitStack() as stack: _ = [stack.enter_context(_.bind(program)) for _ in attrs] yield if self.index_id != -1: GL.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, 0) GL.glBindVertexArray(0) @traitlets.observe("indices") def _set_indices(self, change): arr = change["new"] self.index_id = GL.glGenBuffers(1) GL.glBindBuffer(GL.GL_ELEMENT_ARRAY_BUFFER, self.index_id) GL.glBufferData(GL.GL_ELEMENT_ARRAY_BUFFER, arr.nbytes, arr, GL.GL_STATIC_DRAW)
class VertexAttribute(traitlets.HasTraits): name = traitlets.CUnicode("attr") id = traitlets.CInt(-1) data = traittypes.Array(None, allow_none=True) each = traitlets.CInt(-1) opengl_type = traitlets.CInt(GL.GL_FLOAT) divisor = traitlets.CInt(0) @traitlets.default("id") def _id_default(self): return GL.glGenBuffers(1) @contextmanager def bind(self, program=None): loc = -1 if program is not None: loc = GL.glGetAttribLocation(program.program, self.name) if loc >= 0: GL.glVertexAttribDivisor(loc, self.divisor) _ = GL.glEnableVertexAttribArray(loc) _ = GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.id) if loc >= 0: GL.glVertexAttribPointer(loc, self.each, self.opengl_type, False, 0, None) yield if loc >= 0: GL.glDisableVertexAttribArray(loc) GL.glBindBuffer(GL.GL_ARRAY_BUFFER, 0) @traitlets.observe("data") def _set_data(self, change): arr = change["new"] self.each = arr.shape[-1] self.opengl_type = np_to_gl[arr.dtype.name] with self.bind(): GL.glBufferData(GL.GL_ARRAY_BUFFER, arr.nbytes, arr, GL.GL_STATIC_DRAW)
class Texture(traitlets.HasTraits): texture_name = traitlets.CInt(-1) data = traittypes.Array(None, allow_none=True) channels = GLValue("r32f") min_filter = GLValue("linear") mag_filter = GLValue("linear") @traitlets.default("texture_name") def _default_texture_name(self): return GL.glGenTextures(1) @contextmanager def bind(self, target=0): _ = GL.glActiveTexture(TEX_TARGETS[target]) _ = GL.glBindTexture(self.dim_enum, self.texture_name) yield _ = GL.glActiveTexture(TEX_TARGETS[target]) GL.glBindTexture(self.dim_enum, 0)
class ImageCanvas(ipywidgets.DOMWidget): """ This widget is capable of displaying numpy array as images and draw polygons on them. It scales the images to fit in the view and ensures the polygons are scaled accordingly as well. It is also capable to provide hover/click statuses for the displayed polygons. Args: width (Integer): Width of the widget in pixels; Default **750** height (Integer): Height of the widget in pixels; Default **500** enable_rect (Boolean): Whether to enable the rectangle functionality; Default **True** auto_clear (Boolean): Whether to clear the polygons when drawing a new image; Default **True** enlarge (Boolean): Whether to enlarge an image to take up the most space in the canvas; Default **True** color (String): Default color to draw polygons; Default **#1F77B4** alpha (String): Default alpha fill value for the polygons; Default **00** size (Integer): Default border thickness for the polygons; Default **2** hover_style (Dict): Default hover style (can contain color,alpha and/or size properties); Default **None** click_style (Dict): Default click style (can contain color,alpha and/or size properties); Default **None** Attributes: These attributes can be set and read from your python code to influence the canvas. image (numpy.ndarray): Image data in HWC order. See _validate_image for more information polygons (dict): polygons to draw. See _validate_polygons for more information clicked (Integer): Index of the clicked rectangle hovered (Integer): Index of the hovered rectangle save (Bool): Save image and polygons """ _model_module = traitlets.Unicode('ibb').tag(sync=True) _model_name = traitlets.Unicode('ImageCanvasModel').tag(sync=True) _model_module_version = traitlets.Unicode('2.0.0').tag(sync=True) _view_module = traitlets.Unicode('ibb').tag(sync=True) _view_name = traitlets.Unicode('ImageCanvasView').tag(sync=True) _view_module_version = traitlets.Unicode('2.0.0').tag(sync=True) # Settings width = traitlets.Int(750).tag(sync=True) height = traitlets.Int(500).tag(sync=True) enable_poly = traitlets.Bool(True).tag(sync=True) auto_clear = traitlets.Bool(True) enlarge = traitlets.Bool(True).tag(sync=True) color = traitlets.Unicode('#1F77B4').tag(sync=True) alpha = traitlets.Unicode('00').tag(sync=True) size = traitlets.Int(2).tag(sync=True) hover_style = traitlets.Dict(default_value=None, allow_none=True).tag(sync=True) click_style = traitlets.Dict(default_value=None, allow_none=True).tag(sync=True) # Attributes image = traittypes.Array(None, allow_none=True).tag(sync=True, to_json=array_to_binary) polygons = traitlets.List(None, allow_none=True).tag(sync=True) clicked = traitlets.Int(None, allow_none=True).tag(sync=True) hovered = traitlets.Int(None, allow_none=True).tag(sync=True) save = traitlets.Bool(False).tag(sync=True) def __init__(self, *args, **kwargs): for attr in ('color', 'alpha', 'size'): if attr in kwargs: kwargs[attr] = getattr(self, f'_validate_{attr}')({ 'value': kwargs[attr] }) if 'hover_style' in kwargs and kwargs['hover_style'] != None: try: kwargs['hover_style'] = self._validate_fx( kwargs['hover_style']) except Exception as err: raise ValueError(f'Wrong value in hover_style: {err}') from err if 'click_style' in kwargs and kwargs['click_style'] != None: try: kwargs['click_style'] = self._validate_fx( kwargs['click_style']) except Exception as err: raise ValueError(f'Wrong value in click_style: {err}') from err super().__init__(*args, **kwargs) @traitlets.validate('image') def _validate_image(self, proposal): """ Validate correct image shape and dtype and cast to RGBA uint8 (0-255) Valid data types: - RGBA uint8 (0-255) - RGB uint8 (0-255) - Grayscale uint8 (0-255) - RGBA float (0-1) - RGB float (0-1) - Grayscale float (0-1) """ img = proposal['value'] if self.auto_clear: self.polygons = None self.hovered = None self.clicked = None if img is None: return img if not isinstance(img, np.ndarray): raise TypeError( f'image should by a numpy array or None [{type(img)}]') if img.dtype == np.uint8: if img.ndim == 2: img = np.dstack( (img, img, img, 255 * np.ones(img.shape, dtype=np.uint8))) return img elif img.ndim == 3 and img.shape[2] in (3, 4): if img.shape[2] == 3: img = np.dstack( (img, 255 * np.ones(img.shape[:-1], dtype=np.uint8))) return img else: raise ValueError( f'Image shape not supported [{img.shape}, {img.dtype}]') elif img.dtype in (np.float32, np.float64): if img.ndim == 2: img = np.dstack( (img, img, img, np.ones(img.shape, dtype=img.dtype))) elif img.ndim == 3 and img.shape[2] == 3: img = np.dstack((img, np.ones(img.shape[:-1], dtype=img.dtype))) if img.ndim == 3 and img.shape[2] == 4: return (img * 255).astype(np.uint8) else: raise ValueError( f'Image shape not supported [{img.shape}, {img.dtype}]') else: raise TypeError(f'Image type not supported [{img.dtype}]') @traitlets.validate('polygons') def _validate_polygons(self, proposal): """ Validate correct polygon data Valid data types: - list of dictionaries with keys: coords, color<optional>, alpha<optional>, size<optional> - None (clears) Warning: The individual data values are not validated, as that would slow down everything! It is up to the user to ensure that the values have the following types: - coords: 2D numpy array with X,Y coordinates - color: RGB string - alpha: 2 character HEX string (00-FF) - size: Integer """ poly = proposal['value'] self.hovered = None self.clicked = None if poly is None: return poly if isinstance(poly, list): for p in poly: if ('coords' not in p) or (not isinstance(p['coords'], list)): raise ValueError( 'polygon coords attribute not a np.ndarray or missing') return poly else: raise TypeError(f'Polygons should be a list<dict> [{type(poly)}]') @traitlets.validate('color') def _validate_color(self, proposal): """ Validate correct color type Valid data types: - RGB String: '#XXXXXX' or 'rgb(xx, xx, xx)' (any JS compatible RGB string) """ col = proposal['value'] if not isinstance(col, str): raise TypeError(f'Color should be an RGB string [{type(col)}]') return col @traitlets.validate('alpha') def _validate_alpha(self, proposal): """ Validate default fill alpha Valid types: - String: Hex alpha value (00 - ff) - Integer: integer alpha (0-255) - Float: percentage alpha (0-1) """ return cast_alpha(proposal['value']) @traitlets.validate('size') def _validate_size(self, proposal): """ Validate default border size """ size = proposal['value'] if size < 0: raise ValueError( 'Border size should be bigger or equal than zero [{size}]') return size def _validate_fx(self, val): """ Validate hover/clicked attributes. These attributes should be dicts that can contain color, alpha and/or size values """ if 'color' in val: val['color'] = self._validate_color({'value': val['color']}) else: val['color'] = None if 'alpha' in val: val['alpha'] = self._validate_alpha({'value': val['alpha']}) else: val['alpha'] = None if 'size' in val: val['size'] = self._validate_size({'value': val['size']}) else: val['size'] = None return val
class BaseCamera(traitlets.HasTraits): """Camera object used in the Interactive Data Visualization Parameters ---------- position : :obj:`!iterable`, or 3 element array in code_length The initial position of the camera. focus : :obj:`!iterable`, or 3 element array in code_length A point in space that the camera is looking at. up : :obj:`!iterable`, or 3 element array in code_length The 'up' direction for the camera. fov : float, optional An angle defining field of view in degrees. near_plane : float, optional The distance to the near plane of the perspective camera. far_plane : float, optional The distance to the far plane of the perspective camera. aspect_ratio: float, optional The ratio between the height and the width of the camera's fov. """ # We have to be careful about some of these, as it's possible in-place # operations won't trigger our observation. position = YTPositionTrait([0.0, 0.0, 1.0]) focus = YTPositionTrait([0.0, 0.0, 0.0]) up = traittypes.Array(np.array([0.0, 0.0, 1.0])).valid(ndarray_shape(3), ndarray_ro()) fov = traitlets.Float(45.0) near_plane = traitlets.Float(0.001) far_plane = traitlets.Float(20.0) aspect_ratio = traitlets.Float(8.0 / 6.0) projection_matrix = traittypes.Array(np.zeros( (4, 4))).valid(ndarray_shape(4, 4), ndarray_ro()) view_matrix = traittypes.Array(np.zeros( (4, 4))).valid(ndarray_shape(4, 4), ndarray_ro()) orientation = traittypes.Array(np.zeros(4)).valid(ndarray_shape(4), ndarray_ro()) held = traitlets.Bool(False) @contextlib.contextmanager def hold_traits(self, func): # for some reason, hold_trait_notifications doesn't seem to work here. # So, we use this to block. We also do not want to pass the # notifications once completed. if not self.held: self.held = True func() self.held = False yield @traitlets.default("up") def _default_up(self): return np.array([0.0, 1.0, 0.0]) @traitlets.observe( "position", "focus", "up", "fov", "near_plane", "far_plane", "aspect_ratio", "orientation", ) def compute_matrices(self, change=None): """Regenerate all position, view and projection matrices of the camera.""" with self.hold_traits(self._compute_matrices): pass def update_orientation(self, start_x, start_y, end_x, end_y): """Change camera orientation matrix using delta of mouse's cursor position Parameters ---------- start_x : float initial cursor position in x direction start_y : float initial cursor position in y direction end_x : float final cursor position in x direction end_y : float final cursor position in y direction """ pass def _set_uniforms(self, scene, shader_program): GL.glDepthRange(0.0, 1.0) # Not the same as near/far plane shader_program._set_uniform("projection", self.projection_matrix) shader_program._set_uniform("modelview", self.view_matrix) shader_program._set_uniform( "viewport", np.array(GL.glGetIntegerv(GL.GL_VIEWPORT), dtype="f4")) shader_program._set_uniform("near_plane", self.near_plane) shader_program._set_uniform("far_plane", self.far_plane) shader_program._set_uniform("camera_pos", self.position)
class FrameLinkedPlot(widgets.Box): ploty = tl.List(tl.Union([tl.Integer(), tl.Unicode()]), default_value=[0]) plotx = tl.Union([tl.Integer(), tl.Unicode()], default_value=1) stride = tl.Integer(default_value=1) title = tl.Unicode(default_value='') colvars = tl.List(trait=tl.Union([tt.DataFrame(), tt.Array()]), read_only=True) bokeh = tl.Instance(klass=BokehWrapper, read_only=True) views = tl.List(trait=NGLViewTrait(), read_only=True) sources = tl.List(trait=CDSTrait(), read_only=True) _reuse_colvars = tl.Bool(read_only=True) def_button_layout = widgets.Layout(height='30px') def_view_layout = dict(width='250px', height='300px', flex='0 0 auto') def_figure_layout = widgets.Layout(width='400', height='550') def_box_layout = widgets.Layout(display='flex', flex_flow='column wrap', height='650px', width='100%', align_items='flex-start') def __init__(self, colvars, mdtraj_trajectories, plotx_label=None, ploty_label=None, ploty=None, legendlocation='top_right', **kwargs): super().__init__(**kwargs) # colvars and mdtraj_trajectories can be specified either as a list or a single item if isinstance(mdtraj_trajectories, Trajectory): mdtrajs = [mdtraj_trajectories] else: mdtrajs = list(mdtraj_trajectories) if isinstance(colvars, DataFrame) or isinstance(colvars, Mapping): cvs = [colvars] else: cvs = list(colvars) if len(mdtrajs) != 1 and len(cvs) != len(mdtrajs) and len(ploty) != len(mdtrajs): print(len(mdtrajs), len(cvs), len(ploty)) raise ValueError("Should have either 1 mdtraj trajectory, or one for every colvar") # ploty wasn't given to super, so it's currently the default per the traitlet # So if it was user-specified, we need to make sure its a list and then # throw it in if isinstance(ploty, str): print("ploty is a str") self.ploty = [ploty] * len(cvs) elif ploty is not None: try: self.ploty = list(iter(ploty)) except TypeError: self.ploty = [ploty] * len(cvs) else: if len(cvs) == 1: cvs = cvs * len(self.ploty) if len(cvs) != len(self.ploty): raise ValueError("Couldn't broadcast colvars to plotys: {}, {}".format(len(cvs), len(ploty))) self.set_trait('colvars', cvs) self.layout = copy(self.def_box_layout) plotx = self.plotx ploty = self.ploty self._vlines = {} TOOLS="crosshair,pan,zoom_in,zoom_out,box_zoom,box_select,undo,reset," hover = HoverTool(tooltips=[ ("x", "@x (@plotx_label)"), ("y", "@y (@ploty_label)"), ("time", "@time ns"), ("run", "@run") ]) if plotx_label is None: plotx_label = "Collective variable {}".format(plotx+1) if plotx_label is None: ploty_label = "Collective variable {}".format(ploty+1) p = figure( title=self.title, x_axis_label=plotx_label, y_axis_label=ploty_label, tools=TOOLS) p.add_tools(hover) figure_layout = copy(self.def_figure_layout) bokeh = BokehWrapper(p, layout=figure_layout, showing=False) self.set_trait('bokeh', bokeh) sources, views = self._init_plots_views(mdtrajs, plotx_label, ploty_label) p.legend.click_policy = "hide" p.legend.location = legendlocation if len(cvs) == 1 or len(cvs) == len(mdtrajs): p.legend.visible = False self.set_trait('views', views) self.set_trait('sources', sources) button = widgets.Button( description='Next selected frame', disabled=False, button_style='', # 'success', 'info', 'warning', 'danger' or '' tooltip='Set NGLView frame to the next selected point', layout=copy(self.def_button_layout) ) button.on_click(self._on_button_clicked) self.children =tuple(self.views + [self.bokeh, button]) self.bokeh.show() def _on_button_clicked(self, b): for view in self.views: view.next_selected() def _init_plots_views(self, mdtraj_trajectories, plotx_label, ploty_label): colvars = self.colvars stride = self.stride plotx = self.plotx plotys = self.ploty p = self.figure n = len(colvars) if n >= 3: palette = colorblind['Colorblind'][n] else: palette = colorblind['Colorblind'][3] palette = palette[::-1] view_layout = widgets.Layout(**self.def_view_layout) if len(mdtraj_trajectories) == 1: view_layout.width='500px' view_layout.height='600px' sources = [] views = [] for n,(colvar, traj, ploty) in enumerate(zip_longest(colvars, mdtraj_trajectories, plotys)): if traj is not None: working_traj = traj times = working_traj.time[::stride] if len(colvar) != len(working_traj): raise ValueError("Colvar and trajectory should have same number of frames") if isinstance(colvar, np.ndarray): x = colvar[::stride, plotx] y = colvar[::stride, ploty] else: x = colvar[plotx][::stride] y = colvar[ploty][::stride] if isinstance(ploty, str): this_ploty_label = ploty elif ploty_label is not None: this_ploty_label = ploty_label else: this_ploty_label = str(ploty) ploty_label_list = [this_ploty_label] * len(y) source = ColumnDataSource(data={ 'run': [n]*len(x), 'plotx_label': [plotx_label]*len(x), 'ploty_label': ploty_label_list, 'time': times/1000, 'x': x, 'y': y, 'alphas': [(t)/(times[-1]) for t in times] }) sources.append(source) colour = palette[n-1] if traj is not None: view = show_mdtraj(traj[::stride], gui=False) view.observe(self._update_frame, names=['frame']), view._colour = colour if len(traj.top.select('protein')): view.clear_representations() view.add_cartoon(selection='polymer', color=colour) # view.frame_stride = stride view.layout = view_layout view._set_sync_camera() views.append(view) vline = Span( location=view.frame * traj.timestep * stride + traj.time[0], dimension='height', line_color=colour ) self._vlines[view] = vline p.add_layout(vline) view.link_to_bokeh_ds(source) p.scatter(x='x', y='y', source=source, color=colour, fill_alpha='alphas', legend=this_ploty_label) return sources, views @property def figure(self): return self.bokeh.figure @property def mdtraj_trajectories(self): return (view.original_trajectory for view in self.views) @property def traj(self): traj_list = list(self.mdtraj_trajectories) first_traj = trajlist.pop() for traj in traj_list: if traj is not first_traj: raise ValueError("FrameLinkedPlot has multiple trajectories") return first_traj @property def view(self): if len(self.views) == 1: return self.views[0] else: raise ValueError("FrameLinkedPlot has multiple views") @property def source(self): if len(self.views) == 1: return self.sources[0] else: raise ValueError("FrameLinkedPlot has multiple sources") def _update_frame(self, changes): # selections = [deepcopy(source.selected['1d']) for source in self.sources] view = changes['owner'] new_frame = changes['new'] traj = view.original_trajectory vline = self._vlines[view] vline.location = traj.timestep * new_frame + traj.time[0] # for sele, src in zip(selections, self.sources): # src.selected['1d'] = sele self.bokeh.refresh()
class Traits(traitlets.HasTraits): int = traitlets.Int() float = traitlets.Float() complex = traitlets.Complex() bool = traitlets.Bool() array = traittypes.Array(default_value=(0, 0), dtype=np.float64)
class Model(traitlets.HasTraits): origin = traittypes.Array(None, allow_none=True).valid( check_shape(3), check_dtype("f4") ) vertices = traittypes.Array(None, allow_none=True).valid( check_shape(None, 3), check_dtype("f4") ) indices = traittypes.Array(None, allow_none=True).valid( check_shape(None, 3), check_dtype("i4") ) attributes = traittypes.Array(None, allow_none=True) triangles = traittypes.Array(None, allow_none=True).valid( check_shape(None, 3, 3), check_dtype("f4") ) @classmethod def from_ply(cls, filename): # This is probably not the absolute best way to do this. plydata = PlyData.read(filename) vertices = plydata["vertex"][:] faces = plydata["face"][:] triangles = [] xyz_faces = [] for face in _ensure_triangulated(faces): indices = face[0] vert = vertices[indices] triangles.append(np.array([vert["x"], vert["y"], vert["z"]])) xyz_faces.append(indices) xyz_vert = np.stack([vertices[ax] for ax in "xyz"], axis=-1) xyz_faces = np.stack(xyz_faces) colors = None if "diffuse_red" in vertices.dtype.names: colors = np.stack( [vertices["diffuse_{}".format(c)] for c in ("red", "green", "blue")], axis=-1, ) triangles = np.array(triangles).swapaxes(1, 2) obj = cls( vertices=xyz_vert, indices=xyz_faces.astype('i4'), attributes=colors, triangles=triangles, ) return obj @property def geometry(self): attributes = dict( position=pythreejs.BufferAttribute(self.vertices, normalized=False), index=pythreejs.BufferAttribute( self.indices.ravel(order="C").astype("u4"), normalized=False ), ) if self.attributes is not None: attributes["color"] = pythreejs.BufferAttribute(self.attributes) # Face colors requires # https://speakerdeck.com/yomotsu/low-level-apis-using-three-dot-js?slide=22 # and # https://github.com/mrdoob/three.js/blob/master/src/renderers/shaders/ShaderLib.js geometry = pythreejs.BufferGeometry(attributes=attributes) geometry.exec_three_obj_method("computeFaceNormals") return geometry @cached_property def normals(self): r"""Array of the normal vectors for the triangles in this model.""" v10 = self.triangles[:, 1, :] - self.triangles[:, 0, :] v20 = self.triangles[:, 2, :] - self.triangles[:, 0, :] return np.cross(v10, v20) @cached_property def areas(self): r"""Array of areas for the triangles in this model.""" return 0.5 * np.linalg.norm(self.normals, axis=1) def translate(self, delta): self.vertices = self.vertices + delta def rotate(self, q, origin="barycentric"): """ This expects a quaternion as input. """ pass
class BCC(tl.HasTraits): """BCC lattice maps voronoi cells to bin ids DOES THE GEOMETRIC WINDOW DEFINE THE GRID, OR 1:N contiguous option 1: all idx in range(0, len(binner)) valid option 2: all points in [lower, upper] have valid indices ARE THESE MUTUALLY EXCLUSIVE???? """ ndim = PositiveInt() sizes = tt.Array(dtype='u8') # (upper-lower)/width nside = tt.Array(dtype='u8') # 1 + sizes + 1 (for boundary) lower = tt.Array(dtype='f8') upper = tt.Array(dtype='f8') width = tt.Array(dtype='f8') def get_bin_index(self, value): indices, parity = self._get_bin_indices(value) indices += 1 index = np.sum(self._nside_prod() * indices, axis=-1) return np.left_shift(index, 1) + parity def get_bin_center(self, index): index = np.asarray(index, dtype='i8') parity = np.bitwise_and(index, 1) indices = np.mod( np.right_shift(index, 1)[..., np.newaxis] // self._nside_prod(), self.nside) indices -= 1 return self._get_bin_center_from_indices(indices, parity) def __len__(self): return np.prod(self.nside) if 1: # 'HIDDEN METHODS' def _get_bin_indices(self, value): value = np.asarray(value) tmp = value.copy() tmp -= self.lower tmp /= self.width indices = tmp.astype('i8') tmp -= indices + 0.5 corner_indices = indices - (tmp < 0) parity = 0.25 * self.ndim < np.sum(np.sign(tmp) * tmp, axis=-1) indices = np.where(parity[..., np.newaxis], corner_indices, indices) return indices, parity def _nside_prod(self): _nside_prod = np.ones(self.ndim, dtype='i8') _nside_prod[1:] = np.cumprod(self.nside[:-1]) return _nside_prod def _get_bin_center_from_indices(self, indices, parity): cen = (self.lower + self.width / 2 + self.width * indices + np.where(parity[..., np.newaxis], self.width / 2.0, 0.0)) return cen @tl.default('ndim') def _default_ndim(self): if len(self.sizes.shape) is not 1 or self.sizes.shape[0] < 3: raise tl.TraitError('sizes must have shape (ndim>=3,)') return self.sizes.shape[0] @tl.default('sizes') def _default_sizes(self): raise tl.TraitError('sizes must be specified') @tl.default('lower') def _default_lower(self): return np.zeros((self.ndim, )) @tl.default('upper') def _default_upper(self): return np.ones((self.ndim, )) @tl.default('width') def _default_width(self): return (self.upper - self.lower) / self.sizes @tl.default('nside') def _default_nside(self): return self.sizes + 2 @tl.validate('sizes') def _validate_sizes(self, proposal): return proposal['value'] @tl.validate('nside') def _validate_nside(self, proposal): if not np.all(proposal['value'] == self.sizes + 2): raise tl.TraitError('bad nside, should be sizes + 2') return proposal['value'] @tl.validate('lower') def _validate_lower(self, proposal): lower = proposal['value'] if (self.sizes.shape != lower.shape): err = 'sizes and lower must have same shape. ' err += 'sizes.shape: %s, lower.shape: %s' % (self.sizes.shape, lower.shape) raise tl.TraitError(err) with self.hold_trait_notifications(): if np.any(lower >= self.upper): raise tl.TraitError('lower >= upper!') return lower @tl.validate('upper') def _validate_upper(self, proposal): upper = proposal['value'] if (self.sizes.shape != upper.shape): err = 'sizes and upper must have same shape. ' err += 'sizes.shape: %s, upper.shape: %s' % (self.sizes.shape, upper.shape) raise tl.TraitError(err) with self.hold_trait_notifications(): if np.any(self.lower >= upper): raise tl.TraitError('lower >= upper!') return upper
class Scene(traitlets.HasTraits): ground = traittypes.Array(np.array([0.0, 0.0, 0.0], "f4")).valid(check_dtype("f4"), check_shape(3)) up = traittypes.Array(np.array([0.0, 0.0, 1.0], "f4")).valid(check_dtype("f4"), check_shape(3)) north = traittypes.Array(np.array([0.0, 1.0, 0.0], "f4")).valid(check_dtype("f4"), check_shape(3)) components = traitlets.List(trait=traitlets.Instance(Model)) blasters = traitlets.List(trait=traitlets.Instance(RayBlaster)) meshes = traitlets.List(trait=traitlets.Instance(TriangleMesh)) embree_scene = traitlets.Instance(rtcs.EmbreeScene, args=tuple()) # TODO: Add surface for ground so that reflection from ground # is taken into account def add_component(self, component): self.components = self.components + [component ] # Force traitlet update self.meshes.append(TriangleMesh(self.embree_scene, component.triangles)) def compute_hit_count(self, blaster): output = blaster.compute_count(self) component_counts = {} for ci, component in enumerate(self.components): hits = output["primID"][output["geomID"] == ci] component_counts[ci] = np.bincount( hits[hits >= 0], minlength=component.triangles.shape[0]) return component_counts def get_sun_blaster(self, latitude, longitude, date, direct_ppfd=1.0, diffuse_ppfd=1.0, **kwargs): r"""Get a sun blaster that is adjusted for this scene so that the blaster will never intercept a component in the scene. This distance is determined by computing the maximum distance of any vertex in the scene from the ground parameter. Args: latitude (float): Latitude (in degrees) of the scene. longitude (float): Longitude (in degrees) of the scene. date (datetime.datetime): Time when PPFD should be calculated. This determines the incidence angle of light from the sun. direct_ppfd (float, optional): Direct Photosynthetic Photon Flux Density (PPFD) at the surface of the Earth for the specified location and time. Defaults to 1.0. diffuse_ppfd (float, optional): Diffuse Photosynthetic Photon Flux Density (PPFD) at the surface of the Earth for the specified location and time. Defaults to 1.0. Returns: SunRayBlaster: Blaster tuned to this scene. """ # TODO: Calculate direct/diffuse ppfd from lat/long/date # using pvi if not provided max_distance2 = 0.0 for c in self.components: max_distance2 = max( max_distance2, np.max(np.sum((c.vertices - self.ground)**2, axis=1))) max_distance = np.sqrt(max_distance2) kwargs.setdefault('zenith', self.up * max_distance) kwargs.setdefault('width', 2 * max_distance) kwargs.setdefault('height', 2 * max_distance) kwargs.setdefault('intensity', (direct_ppfd * kwargs['width'] * kwargs['height'])) kwargs.setdefault('diffuse_intensity', diffuse_ppfd) blaster = SunRayBlaster(latitude=latitude, longitude=longitude, date=date, ground=self.ground, north=self.north, **kwargs) return blaster def compute_flux_density(self, light_sources, any_direction=True): r"""Compute the flux density on each scene element from a set of light sources. Values will be calculated from the 'intensity' attribute of the light source blasters such that the flux density will have units of [intensity units] / [distance unit from scene] ** 2. Args: light_sources (list): Set of RayBlasters used to determine the light incident on scene elements. any_direction (bool, optional): If True, light is deposited on component reguardless of if the blaster rays hit the front or back of a component surface. If False, light is only deposited if the blaster rays hit the front. Defaults to True. Returns: dict: Mapping from scene component to an array of flux density values for each triangle in the component. """ if isinstance(light_sources, RayBlaster): light_sources = [light_sources] component_fd = {} for ci, component in enumerate(self.components): component_fd[ci] = np.zeros(component.triangles.shape[0], "f4") for blaster in light_sources: counts = blaster.compute_count(self) any_hits = (counts["primID"] >= 0) for ci, component in enumerate(self.components): idx_hits = np.logical_and(counts["geomID"] == ci, any_hits) norms = component.normals areas = component.areas if isinstance(blaster, OrthographicRayBlaster): component_counts = np.bincount( counts["primID"][idx_hits], minlength=component.triangles.shape[0]) aoi = np.arccos( np.dot(norms, -blaster.forward) / (2.0 * areas * np.linalg.norm(blaster.forward))) if any_direction: aoi[aoi > np.pi / 2] -= np.pi else: aoi[aoi > np.pi / 2] = np.pi # No contribution component_fd[ci] += (component_counts * blaster.ray_intensity * np.cos(aoi) / areas) else: # TODO: This loop can be removed if AOI is calculated # for each intersection by embree (or callback) for idx_ray in np.where(idx_hits)[0]: idx_scene = output["primID"][i] aoi = np.arccos( np.dot(norms[idx_scene], -blaster.directions[idx_ray, :]) / (2.0 * areas[idx_scene] * np.linalg.norm(blaster.directions[idx_ray, :]))) if any_direction: aoi[aoi > np.pi / 2] -= np.pi else: aoi[aoi > np.pi / 2] = np.pi # No contribution component_fd[ci][idx_scene] += (blaster.ray_intensity * np.cos(aoi) / areas[idx_scene]) # Diffuse # TODO: This assumes diffuse light comes from everywhere tilt = np.arccos( np.dot(norms, self.up) / (2.0 * areas * np.linalg.norm(self.up))) component_fd[ci] += pvlib.irradiance.isotropic( np.degrees(tilt), blaster.diffuse_intensity) return component_fd def _ipython_display_(self): # This needs to actually display, which is not the same as returning a display. cam = pythreejs.PerspectiveCamera( position=[25, 35, 100], fov=20, children=[pythreejs.AmbientLight()], ) children = [cam, pythreejs.AmbientLight(color="#dddddd")] material = pythreejs.MeshBasicMaterial(color="#ff0000", vertexColors="VertexColors", side="DoubleSide") for model in self.components: mesh = pythreejs.Mesh(geometry=model.geometry, material=material, position=[0, 0, 0]) children.append(mesh) scene = pythreejs.Scene(children=children) rendererCube = pythreejs.Renderer( camera=cam, background="white", background_opacity=1, scene=scene, controls=[pythreejs.OrbitControls(controlling=cam)], width=800, height=800, ) return rendererCube
class RayBlaster(traitlets.HasTraits): origins = traittypes.Array().valid(check_shape(None, 3), check_dtype("f4")) directions = traittypes.Array().valid(check_shape(None, 3), check_dtype("f4")) intensity = traitlets.CFloat(1.0) diffuse_intensity = traitlets.CFloat(0.0) @property def ray_intensity(self): r"""float: Intensity of single ray.""" return self.intensity / self.origins.shape[0] def cast_once(self, scene, verbose_output=False, query_type=QueryType.DISTANCE): output = scene.embree_scene.run( self.origins, self.directions, query=query_type._value_, output=verbose_output, ) return output def compute_distance(self, scene): output = self.cast_once(scene, verbose_output=False, query_type=QueryType.DISTANCE) return output def compute_count(self, scene): output = self.cast_once(scene, verbose_output=True, query_type=QueryType.INTERSECT) return output def compute_flux_density(self, scene, light_sources, any_direction=True): r"""Compute the flux density on each scene element touched by this blaster from a set of light sources. Args: scene (Scene): Scene to get flux density for. light_sources (list): Set of RayBlasters used to determine the light incident on scene elements. any_direction (bool, optional): If True, the flux is deposited on component reguardless of if the blaster ray hits the front or back of a component surface. If False, flux is only deposited if the blaster ray hits the front. Defaults to True. Returns: array: Total flux density on surfaces intercepted by the rays. """ fd_scene = scene.compute_flux_density(light_sources, any_direction=any_direction) out = np.zeros(self.nx * self.ny, "f4") camera_hits = self.compute_count(scene) for ci, component in enumerate(scene.components): idx_ci = np.where(camera_hits["geomID"] == ci)[0] hits = camera_hits["primID"][idx_ci] out[idx_ci[hits >= 0]] += fd_scene[ci][hits[hits >= 0]] return out
class SunRayBlaster(OrthographicRayBlaster): # ground: Position of center of ray projection on the ground # zenith: Position directly above 'ground' at distance that sun # blaster should be placed. # north: Direction of north on ground from 'ground' latitude = traitlets.Float() longitude = traitlets.Float() date = traitlets.Instance(klass=datetime.datetime) ground = traittypes.Array().valid(check_dtype("f4"), check_shape(3)) zenith = traittypes.Array().valid(check_dtype("f4"), check_shape(3)) north = traittypes.Array().valid(check_dtype("f4"), check_shape(3)) solar_altitude = traitlets.CFloat() solar_azimuth = traitlets.CFloat() solar_distance = traitlets.CFloat() _solpos_info = traittypes.DataFrame() @traitlets.default("_solpos_info") def _solpos_info_default(self): return pvlib.solarposition.get_solarposition(self.date, self.latitude, self.longitude) @traitlets.default("solar_altitude") def _default_solar_altitude(self): solar_altitude = self._solpos_info["apparent_elevation"][0] if solar_altitude < 0: raise ValueError("For the provided lat, long, date, & time " "the sun will be below the horizon.") return solar_altitude @traitlets.default("solar_azimuth") def _default_solar_azimuth(self): return self._solpos_info["azimuth"][0] @property def zenith_direction(self): zd_nonorm = self.zenith - self.ground solar_distance = np.linalg.norm(zd_nonorm) print("sd zd", solar_distance, zd_nonorm) return zd_nonorm / solar_distance @traitlets.default("solar_distance") def _solar_distance_default(self): zd_nonorm = self.zenith - self.ground return np.linalg.norm(zd_nonorm) def solar_rotation(self, point): r"""Rotate a point according to same rotation that moves sun from the zenith to it location in the sky. Args: point (array): 3D point to rotate """ return sun_calc.rotate_u( sun_calc.rotate_u(point, np.radians(90 - self.solar_altitude), self.north), np.radians(90 - self.solar_azimuth), self.zenith_direction) @traitlets.default("forward") def _forward_default(self): # Negative to point from sun to the earth rather than from # eart to the sun return -self.solar_rotation(self.zenith_direction) @traitlets.default("center") def _center_default(self): v = self.ground - self.solar_distance * self.forward offset = max( 0.0, ((self.height / 2.0) - np.abs( np.linalg.norm(v - self.ground) * np.tan(np.radians(self.solar_altitude)))) / 2, ) print(offset) v = v + offset * self.up return v @traitlets.default("up") def _up_default(self): zenith_direction = self.zenith_direction east = np.cross(self.north, zenith_direction) # The "east" used here is not the "east" used elsewhere. # This is the east wrt north etc, but we need an east for blasting from elsewhere. return -self.solar_rotation(east)
class Generator(t.HasTraits): '''Generator Model''' name = t.CUnicode(default_value='GenCo0', help='Name of Generator (str)') generator_bus = t.CUnicode(default_value='Bus0', help='Bus of Generator (str)') generator_voltage = t.CFloat( default_value=1.0, help='Nominal voltage of the generator (p.u.)') base_power = t.CFloat(default_value=100.0, help='Base power of the generator (MVA)') generation_type = t.Enum(['COAL', 'NATURALGAS', 'WIND'], default_value='COAL') minimum_up_time = t.CInt(default_value=0, min=0, help='Minimum up time (hrs)') minimum_down_time = t.CInt(default_value=0, min=0, help='Minimum down time (hrs)') ramp_up_rate = t.CFloat(default_value=0, min=0, help='Ramp up rate (MW/hr)') ramp_down_rate = t.CFloat(default_value=0, min=0, help='Ramp down rate (MW/hr)') maximum_real_power = t.CFloat(default_value=0, min=0, help='Capacity of Generator (MW)') minimum_real_power = t.CFloat(default_value=0, min=0, help='Minimum generation (MW)') maximum_imag_power = t.CFloat(default_value=0, help='Maximum reactive generation (MVAR)') minimum_imag_power = t.CFloat(default_value=0, help='Minimum reactive generation (MVAR)') initial_real_power = t.CFloat(default_value=0, min=0, help='Initial power generation (MW)') initial_imag_power = t.CFloat(default_value=0, min=0, help='Initial power generation (MVAR)') initial_status = t.CBool(default_value=True, min=0, help='Initial status (bool)') startup_time = t.CInt(default_value=0, min=0, help='Startup time (hrs)') shutdown_time = t.CInt(default_value=0, min=0, help='Shutdown time (hrs)') nsegments = t.CInt(default_value=2, min=MINIMUM_COST_CURVE_SEGMENTS, max=MAXIMUM_COST_CURVE_SEGMENTS, help='Number of data points for piecewise linear') cost_curve_points = tt.Array(default_value=[0, 0, 0], minlen=(MINIMUM_COST_CURVE_SEGMENTS + 1), maxlen=(MAXIMUM_COST_CURVE_SEGMENTS + 1)) cost_curve_values = tt.Array(default_value=[0, 0, 0], minlen=(MINIMUM_COST_CURVE_SEGMENTS + 1), maxlen=(MAXIMUM_COST_CURVE_SEGMENTS + 1)) noload_cost = t.CFloat(default_value=0, min=0, help='No-Load Cost of a Generator ($/hr)') startup_cost = t.CFloat(default_value=0, min=0, help='Startup Cost of a Generator ($/hr)') inertia = t.CFloat(allow_none=True, default_value=None, min=0, help='Inertia of generator (NotImplemented)') droop = t.CFloat(allow_none=True, default_value=None, min=0, help='Droop of generator (NotImplemented)') def __init__(self, *args, **kwargs): super(Generator, self).__init__(*args, **kwargs) @property def _npoints(self): return self.nsegments + 1 @property def ramp_rate(self): raise AttributeError( "'{class_name}' object has no attribute 'ramp_rate'. Try 'ramp_up_rate' or 'ramp_down_rate'." .format(class_name=self.__class__.__name__)) @ramp_rate.setter def ramp_rate(self, v): self.ramp_up_rate = v self.ramp_down_rate = v @t.observe('noload_cost') def _callback_noload_cost_update_points_values(self, change): self.cost_curve_values = [change['new']] * self._npoints return change['new'] @t.observe('minimum_real_power') def _callback_minimum_real_power_update_points_values(self, change): self.cost_curve_points = np.linspace(change['new'], self.maximum_real_power, self._npoints) return change['new'] @t.observe('maximum_real_power') def _callback_maximum_real_power_update_points_values(self, change): self.cost_curve_points = np.linspace(self.minimum_real_power, change['new'], self._npoints) self.ramp_rate = self.maximum_real_power return change['new'] @t.observe('nsegments') def _callback_nsegments_update_points_values(self, change): self.cost_curve_points = np.linspace(self.minimum_real_power, self.maximum_real_power, change['new'] + 1) self.cost_curve_values = [self.noload_cost] * (change['new'] + 1) return change['new'] @t.validate('cost_curve_points', 'cost_curve_values') def _validate_max_length(self, proposal): if not len(proposal['value']) == self._npoints: raise t.TraitError( 'len({class_name}().{trait_name}) must be equal to {class_name}().nsegments + 1. Proposed {trait_name} is {proposal}' .format(class_name=proposal['owner'].__class__.__name__, trait_name=proposal['trait'].name, proposal=proposal['value'])) return proposal['value'] @t.validate('ramp_up_rate', 'ramp_down_rate', 'initial_real_power', 'initial_imag_power') def _less_than_maximum_real_power_check(self, proposal): if not proposal['value'] <= self.maximum_real_power: raise t.TraitError( '{class_name}().{trait_name} must be a less than or equal to {class_name}().maximum_real_power.' .format(class_name=proposal['owner'].__class__.__name__, trait_name=proposal['trait'].name)) else: return proposal['value']
class NetworkModel(t.HasTraits): """The NetworkModel object is created from a PSSTCase object, using the PSSTNetwork object. It contains state information neccesary to interactively visualize a network. Parameters ---------- case : PSSTCase An instance of a PSST case. sel_bus : str (Default=None) The name of a bus in the case to focus on. solver : str (Default='glpk') The name of the mixed integer solver you wish to use Attributes ---------- view_buses : :obj:`list` of :obj:`str` The names of the buses to be displayed. pos : pandas.DataFrame All the x,y positions of all the nodes. edges : pandas.DataFrame All the edges between nodes, and their x,y coordiantes. x_edges, y_edges : numpy.Array List of coordinates for 'start' and 'end' node, for all edges in network. bus_x_vals, bus_y_vals : numpy.Array Arrays containing coordinates of buses to be displayed. bus_x_edges, bus_y_edges : numpy.Array List of coordinates for 'start' and 'end' bus for each edge to be displayed. bus_names: :obj:`list` of :obj:`str` List of names of buses to be displayed. gen_x_vals, gen_y_vals : numpy.Array Arrays containing coordinates of generators to be displayed. gen_x_edges, gen_y_edges : numpy.Array List of coordinates for 'start' and 'end' generator for each edge to be displayed. gen_names: :obj:`list` of :obj:`str` List of names of generators to be displayed. load_x_vals, load_y_vals : numpy.Array Arrays containing coordinates of loads to be displayed. load_x_edges, load_y_edges : numpy.Array List of coordinates for 'start' and 'end' bus for each edge to be displayed. load_names: :obj:`list` of :obj:`str` List of names of loads to be displayed. """ case = t.Instance(PSSTCase, help='The original PSSTCase') network = t.Instance(PSSTNetwork) G = t.Instance(nx.Graph) model = t.Instance(PSSTModel, allow_none=True) sel_bus = t.Unicode(allow_none=True) view_buses = t.List(trait=t.Unicode) all_pos = tt.DataFrame(help='DF with all x,y positions of nodes.') all_edges = tt.DataFrame() pos = tt.DataFrame(help='DF with x,y positions only for display nodes') edges = tt.DataFrame() x_edges = tt.Array([]) y_edges = tt.Array([]) bus_x_vals = tt.Array([]) bus_y_vals = tt.Array([]) bus_names = t.List(trait=t.Unicode) bus_x_edges = tt.Array([]) bus_y_edges = tt.Array([]) gen_x_vals = tt.Array([]) gen_y_vals = tt.Array([]) gen_names = t.List(trait=t.Unicode) gen_x_edges = tt.Array([]) gen_y_edges = tt.Array([]) load_x_vals = tt.Array([]) load_y_vals = tt.Array([]) load_names = t.List(trait=t.Unicode) load_x_edges = tt.Array([]) load_y_edges = tt.Array([]) x_min_view = t.CFloat() x_max_view = t.CFloat() y_min_view = t.CFloat() y_max_view = t.CFloat() _VIEW_OFFSET = 50 def __init__(self, case, sel_bus=None, solver='glpk', *args, **kwargs): super(NetworkModel, self).__init__(*args, **kwargs) # Store PPSTCase, PSSTNetwork, and networkx.Graph self.case = case self.network = create_network(case=self.case) self.G = self.network.graph # Try to solve model self.model = build_model(self.case) self.model.solve(solver=solver) # Make full pos DF. self.all_pos = pd.DataFrame(self.network.positions, index=['x', 'y']).T # Make full edges DF, with coordinates and branch index self.all_edges = pd.DataFrame.from_records( [(i, j) for i, j in self.G.edges()], columns=['start', 'end']) # ---> Add coordinates self.all_edges['start_x'] = self.all_edges['start'].map( lambda e: self.all_pos.loc[e]['x']) self.all_edges['end_x'] = self.all_edges['end'].map( lambda e: self.all_pos.loc[e]['x']) self.all_edges['start_y'] = self.all_edges['start'].map( lambda e: self.all_pos.loc[e]['y']) self.all_edges['end_y'] = self.all_edges['end'].map( lambda e: self.all_pos.loc[e]['y']) # ---> Add branch index # new = case.branch[['F_BUS', 'T_BUS']] # new = new.reset_index() # new = new.rename_axis({'F_BUS': 'start', 'T_BUS': 'end', 'index': 'branch_idx'}, axis=1) # self.all_edges = pd.merge(self.all_edges, new, on=['start', 'end'], how='outer') # Make df with all edge data self.x_edges = [ tuple(edge) for edge in self.all_edges[['start_x', 'end_x']].values ] self.y_edges = [ tuple(edge) for edge in self.all_edges[['start_y', 'end_y']].values ] # Set 'start' and 'end' as index for all_edges df # Todo: Refactor so this happens later. self.all_edges.set_index(['start', 'end'], inplace=True) # Set 'sel_bus' (this should in turn set other variables) self.sel_bus = sel_bus