3.3. Traits: building interactive dialogs

author:Didrik Pinte

The Traits project allows you to simply add validation, initialization, delegation, notification and a graphical user interface to Python object attributes.

Truco

In this tutorial we will explore the Traits toolset and learn how to dramatically reduce the amount of boilerplate code you write, do rapid GUI application development, and understand the ideas which underly other parts of the Enthought Tool Suite.

Traits and the Enthought Tool Suite are open source projects licensed under a BSD-style license.

Intended Audience

Intermediate to advanced Python programmers

Requirements

3.3.1. Introduction

Truco

The Enthought Tool Suite enable the construction of sophisticated application frameworks for data analysis, 2D plotting and 3D visualization. These powerful, reusable components are released under liberal BSD-style licenses.

../../_images/ETS.jpg

The main packages of the Enthought Tool Suite are:

  • Traits - component based approach to build our applications.
  • Kiva - 2D primitives supporting path based rendering, affine transforms, alpha blending and more.
  • Enable - object based 2D drawing canvas.
  • Chaco - plotting toolkit for building complex interactive 2D plots.
  • Mayavi - 3D visualization of scientific data based on VTK.
  • Envisage - application plugin framework for building scriptable and extensible applications

In this tutorial, we will focus on Traits.

3.3.2. Example

Throughout this tutorial, we will use an example based on a water resource management simple case. We will try to model a dam and reservoir system. The reservoir and the dams do have a set of parameters :

  • Name
  • Minimal and maximal capacity of the reservoir [hm3]
  • Height and length of the dam [m]
  • Catchment area [km2]
  • Hydraulic head [m]
  • Power of the turbines [MW]
  • Minimal and maximal release [m3/s]
  • Efficiency of the turbines

The reservoir has a known behaviour. One part is related to the energy production based on the water released. A simple formula for approximating electric power production at a hydroelectric plant is P = \rho hrgk, where:

  • P is Power in watts,
  • \rho is the density of water (~1000 kg/m3),
  • h is height in meters,
  • r is flow rate in cubic meters per second,
  • g is acceleration due to gravity of 9.8 m/s2,
  • k is a coefficient of efficiency ranging from 0 to 1.

Truco

Annual electric energy production depends on the available water supply. In some installations the water flow rate can vary by a factor of 10:1 over the course of a year.

The second part of the behaviour is the state of the storage that depends on controlled and uncontrolled parameters :

storage_{t+1} = storage_t + inflows - release - spillage - irrigation

Advertencia

The data used in this tutorial are not real and might even not have sense in the reality.

3.3.3. What are Traits

A trait is a type definition that can be used for normal Python object attributes, giving the attributes some additional characteristics:

  • Standardization:
    • Initialization
    • Validation
    • Deferral
  • Notification

  • Visualization

  • Documentation

A class can freely mix trait-based attributes with normal Python attributes, or can opt to allow the use of only a fixed or open set of trait attributes within the class. Trait attributes defined by a class are automatically inherited by any subclass derived from the class.

The common way of creating a traits class is by extending from the HasTraits base class and defining class traits :

from traits.api import HasTraits, Str, Float
class Reservoir(HasTraits):
name = Str
max_storage = Float

Advertencia

For Traits 3.x users

If using Traits 3.x, you need to adapt the namespace of the traits packages:

  • traits.api should be enthought.traits.api
  • traitsui.api should be enthought.traits.ui.api

Using a traits class like that is as simple as any other Python class. Note that the trait value are passed using keyword arguments:

reservoir = Reservoir(name='Lac de Vouglans', max_storage=605)

3.3.3.1. Initialisation

All the traits do have a default value that initialise the variables. For example, the basic python types do have the following trait equivalents:

Trait Python Type Built-in Default Value
Bool Boolean False
Complex Complex number 0+0j
Float Floating point number 0.0
Int Plain integer 0
Long Long integer 0L
Str String ‘’
Unicode Unicode u’‘

A number of other predefined trait type do exist : Array, Enum, Range, Event, Dict, List, Color, Set, Expression, Code, Callable, Type, Tuple, etc.

Custom default values can be defined in the code:

from traits.api import HasTraits, Str, Float
class Reservoir(HasTraits):
name = Str
max_storage = Float(100)
reservoir = Reservoir(name='Lac de Vouglans')

Complex initialisation

When a complex initialisation is required for a trait, a _XXX_default magic method can be implemented. It will be lazily called when trying to access the XXX trait. For example:

def _name_default(self):
""" Complex initialisation of the reservoir name. """
return 'Undefined'

3.3.3.2. Validation

Every trait does validation when the user tries to set its content:

reservoir = Reservoir(name='Lac de Vouglans', max_storage=605)

reservoir.max_storage = '230'
---------------------------------------------------------------------------
TraitError                                Traceback (most recent call last)
/Users/dpinte/projects/scipy-lecture-notes/advanced/traits/<ipython-input-7-979bdff9974a> in <module>()
----> 1 reservoir.max_storage = '230'

/Users/dpinte/projects/ets/traits/traits/trait_handlers.pyc in error(self, object, name, value)
    166         """
    167         raise TraitError( object, name, self.full_info( object, name, value ),
--> 168                           value )
    169
    170     def arg_error ( self, method, arg_num, object, name, value ):

TraitError: The 'max_storage' trait of a Reservoir instance must be a float, but a value of '23' <type 'str'> was specified.

3.3.3.3. Documentation

By essence, all the traits do provide documentation about the model itself. The declarative approach to the creation of classes makes it self-descriptive:

from traits.api import HasTraits, Str, Float
class Reservoir(HasTraits):
name = Str
max_storage = Float(100)

The desc metadata of the traits can be used to provide a more descriptive information about the trait :

from traits.api import HasTraits, Str, Float
class Reservoir(HasTraits):
name = Str
max_storage = Float(100, desc='Maximal storage [hm3]')

Let’s now define the complete reservoir class:

from traits.api import HasTraits, Str, Float, Range
class Reservoir(HasTraits):
name = Str
max_storage = Float(1e6, desc='Maximal storage [hm3]')
max_release = Float(10, desc='Maximal release [m3/s]')
head = Float(10, desc='Hydraulic head [m]')
efficiency = Range(0, 1.)
def energy_production(self, release):
''' Returns the energy production [Wh] for the given release [m3/s]
'''
power = 1000 * 9.81 * self.head * release * self.efficiency
return power * 3600
if __name__ == '__main__':
reservoir = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 100.0,
head = 60,
efficiency = 0.8
)
release = 80
print 'Releasing {} m3/s produces {} kWh'.format(
release, reservoir.energy_production(release)
)

3.3.3.4. Visualization: opening a dialog

The Traits library is also aware of user interfaces and can pop up a default view for the Reservoir class:

reservoir1 = Reservoir()
reservoir1.edit_traits()
../../_images/reservoir_default_view.png

TraitsUI simplifies the way user interfaces are created. Every trait on a HasTraits class has a default editor that will manage the way the trait is rendered to the screen (e.g. the Range trait is displayed as a slider, etc.).

In the very same vein as the Traits declarative way of creating classes, TraitsUI provides a declarative interface to build user interfaces code:

from traits.api import HasTraits, Str, Float, Range
from traitsui.api import View
class Reservoir(HasTraits):
name = Str
max_storage = Float(1e6, desc='Maximal storage [hm3]')
max_release = Float(10, desc='Maximal release [m3/s]')
head = Float(10, desc='Hydraulic head [m]')
efficiency = Range(0, 1.)
traits_view = View(
'name', 'max_storage', 'max_release', 'head', 'efficiency',
title = 'Reservoir',
resizable = True,
)
def energy_production(self, release):
''' Returns the energy production [Wh] for the given release [m3/s]
'''
power = 1000 * 9.81 * self.head * release * self.efficiency
return power * 3600
if __name__ == '__main__':
reservoir = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 100.0,
head = 60,
efficiency = 0.8
)
reservoir.configure_traits()
../../_images/reservoir_view.png

3.3.3.5. Deferral

Being able to defer the definition of a trait and its value to another object is a powerful feature of Traits.

from traits.api import HasTraits, Instance, DelegatesTo, Float, Range
from reservoir import Reservoir
class ReservoirState(HasTraits):
"""Keeps track of the reservoir state given the initial storage.
"""
reservoir = Instance(Reservoir, ())
min_storage = Float
max_storage = DelegatesTo('reservoir')
min_release = Float
max_release = DelegatesTo('reservoir')
# state attributes
storage = Range(low='min_storage', high='max_storage')
# control attributes
inflows = Float(desc='Inflows [hm3]')
release = Range(low='min_release', high='max_release')
spillage = Float(desc='Spillage [hm3]')
def print_state(self):
print 'Storage\tRelease\tInflows\tSpillage'
str_format = '\t'.join(['{:7.2f}'for i in range(4)])
print str_format.format(self.storage, self.release, self.inflows,
self.spillage)
print '-' * 79
if __name__ == '__main__':
projectA = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 100.0,
hydraulic_head = 60,
efficiency = 0.8
)
state = ReservoirState(reservoir=projectA, storage=10)
state.release = 90
state.inflows = 0
state.print_state()
print 'How do we update the current storage ?'

A special trait allows to manage events and trigger function calls using the magic _xxxx_fired method:

from traits.api import HasTraits, Instance, DelegatesTo, Float, Range, Event
from reservoir import Reservoir
class ReservoirState(HasTraits):
"""Keeps track of the reservoir state given the initial storage.
For the simplicity of the example, the release is considered in
hm3/timestep and not in m3/s.
"""
reservoir = Instance(Reservoir, ())
min_storage = Float
max_storage = DelegatesTo('reservoir')
min_release = Float
max_release = DelegatesTo('reservoir')
# state attributes
storage = Range(low='min_storage', high='max_storage')
# control attributes
inflows = Float(desc='Inflows [hm3]')
release = Range(low='min_release', high='max_release')
spillage = Float(desc='Spillage [hm3]')
update_storage = Event(desc='Updates the storage to the next time step')
def _update_storage_fired(self):
# update storage state
new_storage = self.storage - self.release + self.inflows
self.storage = min(new_storage, self.max_storage)
overflow = new_storage - self.max_storage
self.spillage = max(overflow, 0)
def print_state(self):
print 'Storage\tRelease\tInflows\tSpillage'
str_format = '\t'.join(['{:7.2f}'for i in range(4)])
print str_format.format(self.storage, self.release, self.inflows,
self.spillage)
print '-' * 79
if __name__ == '__main__':
projectA = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 5.0,
hydraulic_head = 60,
efficiency = 0.8
)
state = ReservoirState(reservoir=projectA, storage=15)
state.release = 5
state.inflows = 0
# release the maximum amount of water during 3 time steps
state.update_storage = True
state.print_state()
state.update_storage = True
state.print_state()
state.update_storage = True
state.print_state()

Dependency between objects can be made automatic using the trait Property. The depends_on attribute expresses the dependency between the property and other traits. When the other traits gets changed, the property is invalidated. Again, Traits uses magic method names for the property :

  • _get_XXX for the getter of the XXX Property trait
  • _set_XXX for the setter of the XXX Property trait
from traits.api import HasTraits, Instance, DelegatesTo, Float, Range
from traits.api import Property
from reservoir import Reservoir
class ReservoirState(HasTraits):
"""Keeps track of the reservoir state given the initial storage.
For the simplicity of the example, the release is considered in
hm3/timestep and not in m3/s.
"""
reservoir = Instance(Reservoir, ())
max_storage = DelegatesTo('reservoir')
min_release = Float
max_release = DelegatesTo('reservoir')
# state attributes
storage = Property(depends_on='inflows, release')
# control attributes
inflows = Float(desc='Inflows [hm3]')
release = Range(low='min_release', high='max_release')
spillage = Property(
desc='Spillage [hm3]', depends_on=['storage', 'inflows', 'release']
)
### Private traits. ######################################################
_storage = Float
### Traits property implementation. ######################################
def _get_storage(self):
new_storage = self._storage - self.release + self.inflows
return min(new_storage, self.max_storage)
def _set_storage(self, storage_value):
self._storage = storage_value
def _get_spillage(self):
new_storage = self._storage - self.release + self.inflows
overflow = new_storage - self.max_storage
return max(overflow, 0)
def print_state(self):
print 'Storage\tRelease\tInflows\tSpillage'
str_format = '\t'.join(['{:7.2f}'for i in range(4)])
print str_format.format(self.storage, self.release, self.inflows,
self.spillage)
print '-' * 79
if __name__ == '__main__':
projectA = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 5,
hydraulic_head = 60,
efficiency = 0.8
)
state = ReservoirState(reservoir=projectA, storage=25)
state.release = 4
state.inflows = 0
state.print_state()

Nota

Caching property

Heavy computation or long running computation might be a problem when accessing a property where the inputs have not changed. The @cached_property decorator can be used to cache the value and only recompute them once invalidated.

Let’s extend the TraitsUI introduction with the ReservoirState example:

from traits.api import HasTraits, Instance, DelegatesTo, Float, Range, Property
from traitsui.api import View, Item, Group, VGroup
from reservoir import Reservoir
class ReservoirState(HasTraits):
"""Keeps track of the reservoir state given the initial storage.
For the simplicity of the example, the release is considered in
hm3/timestep and not in m3/s.
"""
reservoir = Instance(Reservoir, ())
name = DelegatesTo('reservoir')
max_storage = DelegatesTo('reservoir')
max_release = DelegatesTo('reservoir')
min_release = Float
# state attributes
storage = Property(depends_on='inflows, release')
# control attributes
inflows = Float(desc='Inflows [hm3]')
release = Range(low='min_release', high='max_release')
spillage = Property(
desc='Spillage [hm3]', depends_on=['storage', 'inflows', 'release']
)
### Traits view ##########################################################
traits_view = View(
Group(
VGroup(Item('name'), Item('storage'), Item('spillage'),
label = 'State', style = 'readonly'
),
VGroup(Item('inflows'), Item('release'), label='Control'),
)
)
### Private traits. ######################################################
_storage = Float
### Traits property implementation. ######################################
def _get_storage(self):
new_storage = self._storage - self.release + self.inflows
return min(new_storage, self.max_storage)
def _set_storage(self, storage_value):
self._storage = storage_value
def _get_spillage(self):
new_storage = self._storage - self.release + self.inflows
overflow = new_storage - self.max_storage
return max(overflow, 0)
def print_state(self):
print 'Storage\tRelease\tInflows\tSpillage'
str_format = '\t'.join(['{:7.2f}'for i in range(4)])
print str_format.format(self.storage, self.release, self.inflows,
self.spillage)
print '-' * 79
if __name__ == '__main__':
projectA = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 5,
hydraulic_head = 60,
efficiency = 0.8
)
state = ReservoirState(reservoir=projectA, storage=25)
state.release = 4
state.inflows = 0
state.print_state()
state.configure_traits()
../../_images/reservoir_state_view.png

Some use cases need the delegation mechanism to be broken by the user when setting the value of the trait. The PrototypeFrom trait implements this behaviour.

from traits.api import HasTraits, Str, Float, Range, PrototypedFrom, Instance
class Turbine(HasTraits):
turbine_type = Str
power = Float(1.0, desc='Maximal power delivered by the turbine [Mw]')
class Reservoir(HasTraits):
name = Str
max_storage = Float(1e6, desc='Maximal storage [hm3]')
max_release = Float(10, desc='Maximal release [m3/s]')
head = Float(10, desc='Hydraulic head [m]')
efficiency = Range(0, 1.)
turbine = Instance(Turbine)
installed_capacity = PrototypedFrom('turbine', 'power')
if __name__ == '__main__':
turbine = Turbine(turbine_type='type1', power=5.0)
reservoir = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 100.0,
head = 60,
efficiency = 0.8,
turbine = turbine,
)
print 'installed capacity is initialised with turbine.power'
print reservoir.installed_capacity
print '-' * 15
print 'updating the turbine power updates the installed capacity'
turbine.power = 10
print reservoir.installed_capacity
print '-' * 15
print 'setting the installed capacity breaks the link between turbine.power'
print 'and the installed_capacity trait'
reservoir.installed_capacity = 8
print turbine.power, reservoir.installed_capacity

3.3.3.6. Notification

Traits implements a Listener pattern. For each trait a list of static and dynamic listeners can be fed with callbacks. When the trait does change, all the listeners are called.

Static listeners are defined using the _XXX_changed magic methods:

from traits.api import HasTraits, Instance, DelegatesTo, Float, Range
from reservoir import Reservoir
class ReservoirState(HasTraits):
"""Keeps track of the reservoir state given the initial storage.
"""
reservoir = Instance(Reservoir, ())
min_storage = Float
max_storage = DelegatesTo('reservoir')
min_release = Float
max_release = DelegatesTo('reservoir')
# state attributes
storage = Range(low='min_storage', high='max_storage')
# control attributes
inflows = Float(desc='Inflows [hm3]')
release = Range(low='min_release', high='max_release')
spillage = Float(desc='Spillage [hm3]')
def print_state(self):
print 'Storage\tRelease\tInflows\tSpillage'
str_format = '\t'.join(['{:7.2f}'for i in range(4)])
print str_format.format(self.storage, self.release, self.inflows,
self.spillage)
print '-' * 79
### Traits listeners #####################################################
def _release_changed(self, new):
"""When the release is higher than zero, warn all the inhabitants of
the valley.
"""
if new > 0:
print 'Warning, we are releasing {} hm3 of water'.format(new)
if __name__ == '__main__':
projectA = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 100.0,
hydraulic_head = 60,
efficiency = 0.8
)
state = ReservoirState(reservoir=projectA, storage=10)
state.release = 90
state.inflows = 0
state.print_state()

The static trait notification signatures can be:

  • def _release_changed(self):

    pass

  • def _release_changed(self, new):

    pass

  • def _release_changed(self, old, new):

    pass

  • def _release_changed(self, name, old, new

    pass

Listening to all the changes

To listen to all the changes on a HasTraits class, the magic _any_trait_changed method can be implemented.

In many situations, you do not know in advance what type of listeners need to be activated. Traits offers the ability to register listeners on the fly with the dynamic listeners

from reservoir import Reservoir
from reservoir_state_property import ReservoirState
def wake_up_watchman_if_spillage(new_value):
if new_value > 0:
print 'Wake up watchman! Spilling {} hm3'.format(new_value)
if __name__ == '__main__':
projectA = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 100.0,
hydraulic_head = 60,
efficiency = 0.8
)
state = ReservoirState(reservoir=projectA, storage=10)
#register the dynamic listener
state.on_trait_change(wake_up_watchman_if_spillage, name='spillage')
state.release = 90
state.inflows = 0
state.print_state()
print 'Forcing spillage'
state.inflows = 100
state.release = 0
print 'Why do we have two executions of the callback ?'

The dynamic trait notification signatures are not the same as the static ones :

  • def wake_up_watchman():

    pass

  • def wake_up_watchman(new):

    pass

  • def wake_up_watchman(name, new):

    pass

  • def wake_up_watchman(object, name, new):

    pass

  • def wake_up_watchman(object, name, old, new):

    pass

Removing a dynamic listener can be done by:
  • calling the remove_trait_listener method on the trait with the listener method as argument,
  • calling the on_trait_change method with listener method and the keyword remove=True,
  • deleting the instance that holds the listener.

Listeners can also be added to classes using the on_trait_change decorator:

from traits.api import HasTraits, Instance, DelegatesTo, Float, Range
from traits.api import Property, on_trait_change
from reservoir import Reservoir
class ReservoirState(HasTraits):
"""Keeps track of the reservoir state given the initial storage.
For the simplicity of the example, the release is considered in
hm3/timestep and not in m3/s.
"""
reservoir = Instance(Reservoir, ())
max_storage = DelegatesTo('reservoir')
min_release = Float
max_release = DelegatesTo('reservoir')
# state attributes
storage = Property(depends_on='inflows, release')
# control attributes
inflows = Float(desc='Inflows [hm3]')
release = Range(low='min_release', high='max_release')
spillage = Property(
desc='Spillage [hm3]', depends_on=['storage', 'inflows', 'release']
)
### Private traits. ######################################################
_storage = Float
### Traits property implementation. ######################################
def _get_storage(self):
new_storage = self._storage - self.release + self.inflows
return min(new_storage, self.max_storage)
def _set_storage(self, storage_value):
self._storage = storage_value
def _get_spillage(self):
new_storage = self._storage - self.release + self.inflows
overflow = new_storage - self.max_storage
return max(overflow, 0)
@on_trait_change('storage')
def print_state(self):
print 'Storage\tRelease\tInflows\tSpillage'
str_format = '\t'.join(['{:7.2f}'for i in range(4)])
print str_format.format(self.storage, self.release, self.inflows,
self.spillage)
print '-' * 79
if __name__ == '__main__':
projectA = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 5,
hydraulic_head = 60,
efficiency = 0.8
)
state = ReservoirState(reservoir=projectA, storage=25)
state.release = 4
state.inflows = 0

The patterns supported by the on_trait_change method and decorator are powerful. The reader should look at the docstring of HasTraits.on_trait_change for the details.

3.3.3.7. Some more advanced traits

The following example demonstrate the usage of the Enum and List traits :

from traits.api import HasTraits, Str, Float, Range, Enum, List
from traitsui.api import View, Item
class IrrigationArea(HasTraits):
name = Str
surface = Float(desc='Surface [ha]')
crop = Enum('Alfalfa', 'Wheat', 'Cotton')
class Reservoir(HasTraits):
name = Str
max_storage = Float(1e6, desc='Maximal storage [hm3]')
max_release = Float(10, desc='Maximal release [m3/s]')
head = Float(10, desc='Hydraulic head [m]')
efficiency = Range(0, 1.)
irrigated_areas = List(IrrigationArea)
def energy_production(self, release):
''' Returns the energy production [Wh] for the given release [m3/s]
'''
power = 1000 * 9.81 * self.head * release * self.efficiency
return power * 3600
traits_view = View(
Item('name'),
Item('max_storage'),
Item('max_release'),
Item('head'),
Item('efficiency'),
Item('irrigated_areas'),
resizable = True
)
if __name__ == '__main__':
upper_block = IrrigationArea(name='Section C', surface=2000, crop='Wheat')
reservoir = Reservoir(
name='Project A',
max_storage=30,
max_release=100.0,
head=60,
efficiency=0.8,
irrigated_areas=[upper_block]
)
release = 80
print 'Releasing {} m3/s produces {} kWh'.format(
release, reservoir.energy_production(release)
)

Trait listeners can be used to listen to changes in the content of the list to e.g. keep track of the total crop surface on linked to a given reservoir.

from traits.api import HasTraits, Str, Float, Range, Enum, List, Property
from traitsui.api import View, Item
class IrrigationArea(HasTraits):
name = Str
surface = Float(desc='Surface [ha]')
crop = Enum('Alfalfa', 'Wheat', 'Cotton')
class Reservoir(HasTraits):
name = Str
max_storage = Float(1e6, desc='Maximal storage [hm3]')
max_release = Float(10, desc='Maximal release [m3/s]')
head = Float(10, desc='Hydraulic head [m]')
efficiency = Range(0, 1.)
irrigated_areas = List(IrrigationArea)
total_crop_surface = Property(depends_on='irrigated_areas.surface')
def _get_total_crop_surface(self):
return sum([iarea.surface for iarea in self.irrigated_areas])
def energy_production(self, release):
''' Returns the energy production [Wh] for the given release [m3/s]
'''
power = 1000 * 9.81 * self.head * release * self.efficiency
return power * 3600
traits_view = View(
Item('name'),
Item('max_storage'),
Item('max_release'),
Item('head'),
Item('efficiency'),
Item('irrigated_areas'),
Item('total_crop_surface'),
resizable = True
)
if __name__ == '__main__':
upper_block = IrrigationArea(name='Section C', surface=2000, crop='Wheat')
reservoir = Reservoir(
name='Project A',
max_storage=30,
max_release=100.0,
head=60,
efficiency=0.8,
irrigated_areas=[upper_block],
)
release = 80
print 'Releasing {} m3/s produces {} kWh'.format(
release, reservoir.energy_production(release)
)

The next example shows how the Array trait can be used to feed a specialised TraitsUI Item, the ChacoPlotItem:

import numpy as np
from traits.api import HasTraits, Array, Instance, Float, Property
from traits.api import DelegatesTo
from traitsui.api import View, Item, Group
from chaco.chaco_plot_editor import ChacoPlotItem
from reservoir import Reservoir
class ReservoirEvolution(HasTraits):
reservoir = Instance(Reservoir)
name = DelegatesTo('reservoir')
inflows = Array(dtype=np.float64, shape=(None))
releass = Array(dtype=np.float64, shape=(None))
initial_stock = Float
stock = Property(depends_on='inflows, releases, initial_stock')
month = Property(depends_on='stock')
### Traits view ##########################################################
traits_view = View(
Item('name'),
Group(
ChacoPlotItem('month', 'stock', show_label=False),
),
width = 500,
resizable = True
)
### Traits properties ####################################################
def _get_stock(self):
"""
fixme: should handle cases where we go over the max storage
"""
return self.initial_stock + (self.inflows - self.releases).cumsum()
def _get_month(self):
return np.arange(self.stock.size)
if __name__ == '__main__':
reservoir = Reservoir(
name = 'Project A',
max_storage = 30,
max_release = 100.0,
head = 60,
efficiency = 0.8
)
initial_stock = 10.
inflows_ts = np.array([6., 6, 4, 4, 1, 2, 0, 0, 3, 1, 5, 3])
releases_ts = np.array([4., 5, 3, 5, 3, 5, 5, 3, 2, 1, 3, 3])
view = ReservoirEvolution(
reservoir = reservoir,
inflows = inflows_ts,
releases = releases_ts
)
view.configure_traits()
../../_images/reservoir_evolution.png