brancher.py

>>> from minimock import Mock
>>> import numpy as np
>>> import datetime as dt

class Brancher(object):

Simulator branch construction handler.

Attributes: - _branching_groups: dictionary containing the branching groups and branch tasks - _task_index: mapping from branching group and branch task name to task indice, dictionary - _group_index: mapping from branching group name to a deque of task indices - _amap: active branching task array, a boolean numpy array with dimensions (iterations, branches, objects, branch_tasks) - _parent_branches: current parent branch target index in a structure similar to data Handler’s target index - _new_branches: new branch target index in a structure similar to data Handler’s target index - _cur_task: current branching task object - _is_branching: is branching on at the moment, boolean - _aborted: was branching aborted for object, boolean array - _branching_info: branching history container, a dictionary where (iteration, new branch, object) -tuples are keys and (parent branch, operation group, operation name, branch task) -tuples are values.

Branching groups are given in a dictionary, where branching group names are the keys and the values are lists that contain a list of branching_group parts and a boolean indicating whether free branching is allowed. Branching_group parts consist of conditions (in the example below just None values) and branching task dictionaries. The branching task dictionaries themselves have branch task names as keys and lists of branch names as values

>>> lexicon = Mock('Lexicon')
>>> lexicon.is_child = lambda x, y: x > y
>>> bgroups = {
...     'bg1': [[(None, { 'task1': ['branch1', 'branch2'],
...                     'task2': ['branch3']})], False],
...     'bg2': [[(None, {'task3': ['branch4']})], False]}

def __init__(self, iterations, bgroups, no_op_branch, separate_branching_groups, branch_limit):

Initialize a branching handler:

>>> from simo.matrix.brancher import Brancher
>>> br = Brancher(5, bgroups, False, False, None)

def _set_branch_groups(self):

Construct branching task and group indices from the branching group dictionaries:

>>> br._set_branch_groups()
>>> br._task_index['bg1'][('task1','branch1')]
0
>>> br._task_index['bg1'][('task1','branch2')]
1
>>> br._task_index['bg1'][('task2','branch3')]
2
>>> br._task_index['bg2'][('task3','branch4')]
3
>>> br._group_index['bg1']
(0, deque([0, 1, 2]), False)
>>> br._group_index['bg2']
(1, deque([3]), False)
>>> br._active_tasks
array([], shape=(5, 1, 0, 4), dtype=bool)
>>> br._parent_branches
>>> br._new_branches

def add_object(self, branch, num):

Add a number of main level objects to the brancher. This method is needed when constructing the data matrix for the simulation. The method increases active branching task mapping (amap) size in object (3rd) dimension

>>> br.add_object(1)
>>> br._active_tasks[0,0,:,:]
array([[ True,  True,  True,  True]], dtype=bool)

>>> br.add_object(1)
>>> br._active_tasks[0,0,:,:]
array([[ True,  True,  True,  True],
       [ True,  True,  True,  True]], dtype=bool)

def add_branch(self, iteration, frombranch, tobranch, obj):

Add a new branch to the data matrix by copying values from “parent” branch. The method increases active branching task mapping (amap) size in branch (2nd) dimension

>>> br._active_tasks[0,0,1,1] = False
>>> br.add_branch(0, 0, 1, 1)
>>> br._active_tasks[0,1,:,:]
array([[ True,  True,  True,  True],
       [ True, False,  True,  True]], dtype=bool)

def start(self, task, branchname, data, simlevel, depth, pred_mem, oper_mem):

Start branching for objects defined in the target index (tind):

>>> br._active_tasks[1,0,1,1:] = False
>>> br._active_tasks[4,0,0,(2,3)] = False
>>> br._active_bgs[1,0,1,1:] = False
>>> task = Mock('Task')
>>> task.branching_group = 'bg1'
>>> task.name = 'task1'
>>> data = Mock('Data')

Make sure that dummy Handler’s get_tind method returns target index with four objects (iteration, branch, object index, optional data start index, optional data stop index):

>>> data.get_tind.mock_returns = (
...     np.array([[0,0,0,0,0],
...               [1,0,1,0,0],
...               [4,0,0,0,0],
...               [4,0,1,0,0]], dtype=int),
...     None)

Dummy Handler’s add_branch should return the target index for the new branches:

>>> data.add_branch.mock_returns = \
...     np.array([[0,1,0,0,0],
...               [1,2,1,0,0],
...               [4,3,0,0,0]], dtype=int)
>>> oper_mem = Mock('OperationMemory')
>>> pred_mem = Mock('PredictionMemory')

Dummy handler’s linkage should contain the links to the relatives of both the parent and the new branch:

>>> data.linkage.links = {
...     (0,0): {(1,0): {0: [0], 2: [0], 3: [0,1,2,3]}},
...     (1,0): {(1,1): {0: [0], 2: [1], 3: [1,2,3,4]}},
...     (4,0): {(1,0): {0: [0], 2: [2], 3: [2,3,4,5]},
...             (1,1): {0: [0], 2: [3], 3: [3,4,5,6]}},
...     (0,1): {(1,0): {0: [0], 2: [0], 3: [0,1,2,3]}},
...     (1,2): {(1,1): {0: [0], 2: [1], 3: [1,2,3,4]}},
...     (4,3): {(1,0): {0: [0], 2: [3], 3: [3,4,5,6]}}
... }

Start the actual branching:

>>> br.start(task, 'branch1', data, lexicon, 1, 0, pred_mem, oper_mem)
Called Data.get_tind(1, 0)
Called Data.add_branch(
    array([[0, 0, 0, 0, 0],
       [4, 0, 0, 0, 0],
       [4, 0, 1, 0, 0]]))
Called Data.block(
    1,
    0,
    array([[0, 0, 0, 0, 0],
       [4, 0, 0, 0, 0],
       [4, 0, 1, 0, 0]]))
Called PredictionMemory.add_branch(array([0, 0, 0]), array([0, 1, 0]), 1)
Called OperationMemory.add_branch(array([0, 0, 0]), array([0, 1, 0]), 1)
Called PredictionMemory.add_branch((0, 0, 0), (0, 1, 0), 2)
Called OperationMemory.add_branch((0, 0, 0), (0, 1, 0), 2)
Called PredictionMemory.add_branch((0, 0, 0), (0, 1, 0), 3)
Called OperationMemory.add_branch((0, 0, 0), (0, 1, 0), 3)
Called PredictionMemory.add_branch((0, 0, 1), (0, 1, 1), 3)
Called OperationMemory.add_branch((0, 0, 1), (0, 1, 1), 3)
Called PredictionMemory.add_branch((0, 0, 2), (0, 1, 2), 3)
Called OperationMemory.add_branch((0, 0, 2), (0, 1, 2), 3)
Called PredictionMemory.add_branch((0, 0, 3), (0, 1, 3), 3)
Called OperationMemory.add_branch((0, 0, 3), (0, 1, 3), 3)
Called PredictionMemory.add_branch(array([4, 0, 0]), array([1, 2, 1]), 1)
Called OperationMemory.add_branch(array([4, 0, 0]), array([1, 2, 1]), 1)
Called PredictionMemory.add_branch((4, 0, 2), (1, 2, 1), 2)
Called OperationMemory.add_branch((4, 0, 2), (1, 2, 1), 2)
Called PredictionMemory.add_branch((4, 0, 2), (1, 2, 1), 3)
Called OperationMemory.add_branch((4, 0, 2), (1, 2, 1), 3)
Called PredictionMemory.add_branch((4, 0, 3), (1, 2, 2), 3)
Called OperationMemory.add_branch((4, 0, 3), (1, 2, 2), 3)
Called PredictionMemory.add_branch((4, 0, 4), (1, 2, 3), 3)
Called OperationMemory.add_branch((4, 0, 4), (1, 2, 3), 3)
Called PredictionMemory.add_branch((4, 0, 5), (1, 2, 4), 3)
Called OperationMemory.add_branch((4, 0, 5), (1, 2, 4), 3)
Called PredictionMemory.add_branch(array([4, 0, 1]), array([4, 3, 0]), 1)
Called OperationMemory.add_branch(array([4, 0, 1]), array([4, 3, 0]), 1)
Called PredictionMemory.add_branch((4, 0, 3), (4, 3, 3), 2)
Called OperationMemory.add_branch((4, 0, 3), (4, 3, 3), 2)
Called PredictionMemory.add_branch((4, 0, 3), (4, 3, 3), 3)
Called OperationMemory.add_branch((4, 0, 3), (4, 3, 3), 3)
Called PredictionMemory.add_branch((4, 0, 4), (4, 3, 4), 3)
Called OperationMemory.add_branch((4, 0, 4), (4, 3, 4), 3)
Called PredictionMemory.add_branch((4, 0, 5), (4, 3, 5), 3)
Called OperationMemory.add_branch((4, 0, 5), (4, 3, 5), 3)
Called PredictionMemory.add_branch((4, 0, 6), (4, 3, 6), 3)
Called OperationMemory.add_branch((4, 0, 6), (4, 3, 6), 3)
>>> br._parent_branches
array([[0, 0, 0, 0, 0],
       [1, 0, 1, 0, 0],
       [4, 0, 0, 0, 0],
       [4, 0, 1, 0, 0]])
>>> br._new_branches
array([[0, 1, 0, 0, 0],
       [1, 0, 1, 0, 0],
       [1, 2, 1, 0, 0],
       [4, 3, 0, 0, 0]])
>>> br.is_branching
True
>>> br._task  
<Mock ... Task>

Data handler is responsible for adding new branches using brancher’s add_branch method::

>>> br.add_branch(0, 0, 1, 0)
>>> br.add_branch(1, 0, 2, 0)
>>> br.add_branch(4, 0, 3, 0)

def check_operation(self, model, data, level, depth):

Check if the current operation is a branching operation and if a new branch should be created or not. Simulator uses data handler’s get_tind method to deduce if the operation was done to any of the current branching objects.

>>> data.get_tind.mock_returns = (
...     np.array([[0,1,0,0,0]], dtype=int),
...     None)
>>> task.branching_ops = (1,5,6)
>>> model = Mock('Model')
>>> model.group = 'mock group'
>>> model.name = 'mock name'
>>> model.index = 0
>>> model.dropped = np.array([[0,1,0,0,0]], dtype=int)
>>> br.check_operation(model, data, 1, 0)
>>> br._aborted
array([ True,  True,  True,  True], dtype=bool)

>>> model.index = 5
>>> br.check_operation(model, data, 1, 0)
Called Data.get_tind(1, 0)
>>> br._aborted
array([ True,  True,  True,  True], dtype=bool)
>>> model.dropped = None
>>> br.check_operation(model, data, 1, 0)
Called Data.get_tind(1, 0)
>>> br._aborted
array([False,  True,  True,  True], dtype=bool)

>>> data.get_tind.mock_returns = (
...     np.array([[4,3,0,0,0]], dtype=int),
...     None)
>>> br.check_operation(model, data, 1, 0)
Called Data.get_tind(1, 0)
>>> br._aborted
array([False,  True,  True, False], dtype=bool)

def stop(self, branchname, data, timespan, simlevel, depth, pred_mem, oper_mem, lexicon, time_step_level, time_step_ind, opres_queue, opres_before_branching):

Stop branching and reset initialized branching targets and branching task

Parameters

branchname -- new branch name as string
data -- data handler
timespan -- timespan control object
sim_level -- simulation level indice
depth -- model chain depth indice
pred_mem -- prediction model memory structure
oper_mem -- operation model memory structure
time_step_level -- time step variable level, int
time_step_ind -- time step variable indice, int
opres_queue -- queued operation results, list
opres_before_branching -- queued operation results before branching, list

>>> data.get_date.mock_returns = np.array([dt.date(2000,1,1),
...                                        dt.date(2000,1,1),
...                                        dt.date(2000,1,1),
...                                        dt.date(2000,1,1)],
...                                       dtype=dt.date)
>>> data.get_value.mock_returns = (np.array([3.,3.,3.,3.]), None,)
>>> opres_queue = [(None, np.array([[0,0,0,0,0]]), (0, 0, 'stand-1', None)),
...                (None, np.array([[1,0,1,0,0]]), (1, 0, 'stand-1', None))]
>>> opres_before_branching = [(None, np.array([[0,0,0,0,0]]), (0, 0, 'stand-1', None))]
>>> timespan = Mock('Timespan')
>>> timespan.time_step = 3
>>> timespan.unit = 'year'
>>> lexicon = Mock('Lexicon')
>>> lexicon.is_child = (lambda x,y: x > y)
>>> br.stop('branch1', data, timespan, 1, 0, pred_mem, oper_mem, lexicon,
...         1, 15, opres_queue, opres_before_branching)
Called Data.get_date(
    array([[0, 1, 0, 0, 0],
       [1, 0, 1, 0, 0],
       [1, 2, 1, 0, 0],
       [4, 3, 0, 0, 0]]))
Called Data.get_tind(1, 0)
Called Data.get_value(array([[4, 3, 0, 0, 0]]), 15)
Called PredictionMemory.del_objects(1, 2, [1], 1)
Called OperationMemory.del_objects(1, 2, [1], 1)
Called PredictionMemory.del_objects(1, 2, [1], 2)
Called OperationMemory.del_objects(1, 2, [1], 2)
Called PredictionMemory.del_objects(1, 2, [1, 2, 3, 4], 3)
Called OperationMemory.del_objects(1, 2, [1, 2, 3, 4], 3)
Called Data.del_branch(array([[1, 2, 1, 0, 0]]), 0)
Called Data.release(
    1,
    0,
    array([[0, 0, 0, 0, 0],
       [1, 0, 1, 0, 0],
       [4, 0, 0, 0, 0],
       [4, 0, 1, 0, 0]]))
array([[0, 0, 0, 0, 0],
       [4, 0, 1, 0, 0],
       [0, 1, 0, 0, 0],
       [4, 3, 0, 0, 0]])
>>> br._branching_info[(0,1,0)]  
(0, 'mock group', 'mock name', 'branch1', 'bg1', datetime.date(2002, 12, 31))
>>> br._branching_info[(4,3,0)]  
(0, 'mock group', 'mock name', 'branch1', 'bg1', datetime.date(2002, 12, 31))
>>> br.is_branching
False
>>> br._parent_branches
>>> br._new_branches
>>> br._task

Check that in iteration 0 and branch 1, object 0 is associated with branching group 0 and object 1 is not yet associated with any branching group (value -1):

>>> list(br._bg_bind[0,1,:])
[0, -1]

Check that in iteration 4 and branch 3, object 0 is associated with branching group 0 and object 1 is not yet associated with any branching group:

>>> list(br._bg_bind[4,3,:])
[0, -1]

def is_blocked(self, task, branchname):

Check if current branching group / task combination is blocked, i.e. all tasks in the group have been done for all units already

>>> br.is_blocked(task, 'branch1')
False
>>> br._active_tasks[:,:,:,:] = False
>>> br.is_blocked(task, 'branch1')
True

Do a second branching with the same branching group as the previous branch. This should result in new branches for the both objects, as they are not yet associated with any branch:

>>> task = Mock('Task')
>>> task.branching_group = 'bg1'
>>> task.name = 'task2'
>>> data = Mock('Data')

Make sure that dummy Handler’s get_tind method returns target index with two objects (iteration, branch, object index, optional data start index, optional data stop index):

>>> data.get_tind.mock_returns = (
...     np.array([[0,1,0,0,0],
...               [4,3,1,0,0]], dtype=int),
...     None)

Dummy Handler’s add_branch should return the target index for the new branches:

>>> data.add_branch.mock_returns = \
...     np.array([[0,2,0,0,0],
...               [4,4,1,0,0]], dtype=int)
>>> oper_mem = Mock('OperationMemory')
>>> pred_mem = Mock('PredictionMemory')
>>> br._active_tasks[0,1,0,2] = True
>>> br._active_tasks[4,3,1,2] = True
>>> from collections import deque
>>> br._group_index['bg1'] = (0, deque([0, 1, 2]), True)
>>> data.linkage.links = {
...     (0,1): {(1,0): {0: [0], 2: [0], 3: [0,1,2,3]}},
...     (4,3): {(1,1): {0: [0], 2: [1], 3: [1,2,3,4]}},
...     (0,2): {(1,0): {0: [0], 2: [0], 3: [0,1,2,3]}},
...     (4,4): {(1,1): {0: [0], 2: [3], 3: [3,4,5,6]}}
... }

Start the actual branching:

>>> br.start(task, 'branch3', data, lexicon, 1, 0, pred_mem, oper_mem)
Called Data.get_tind(1, 0)
Called Data.add_branch(array([[0, 1, 0, 0, 0],
       [4, 3, 1, 0, 0]]))
Called Data.block(1, 0, array([[0, 1, 0, 0, 0],
       [4, 3, 1, 0, 0]]))
Called PredictionMemory.add_branch(array([0, 1, 0]), array([0, 2, 0]), 1)
Called OperationMemory.add_branch(array([0, 1, 0]), array([0, 2, 0]), 1)
Called PredictionMemory.add_branch((0, 1, 0), (0, 2, 0), 2)
Called OperationMemory.add_branch((0, 1, 0), (0, 2, 0), 2)
Called PredictionMemory.add_branch((0, 1, 0), (0, 2, 0), 3)
Called OperationMemory.add_branch((0, 1, 0), (0, 2, 0), 3)
Called PredictionMemory.add_branch((0, 1, 1), (0, 2, 1), 3)
Called OperationMemory.add_branch((0, 1, 1), (0, 2, 1), 3)
Called PredictionMemory.add_branch((0, 1, 2), (0, 2, 2), 3)
Called OperationMemory.add_branch((0, 1, 2), (0, 2, 2), 3)
Called PredictionMemory.add_branch((0, 1, 3), (0, 2, 3), 3)
Called OperationMemory.add_branch((0, 1, 3), (0, 2, 3), 3)
Called PredictionMemory.add_branch(array([4, 3, 1]), array([4, 4, 1]), 1)
Called OperationMemory.add_branch(array([4, 3, 1]), array([4, 4, 1]), 1)
Called PredictionMemory.add_branch((4, 3, 1), (4, 4, 3), 2)
Called OperationMemory.add_branch((4, 3, 1), (4, 4, 3), 2)
Called PredictionMemory.add_branch((4, 3, 1), (4, 4, 3), 3)
Called OperationMemory.add_branch((4, 3, 1), (4, 4, 3), 3)
Called PredictionMemory.add_branch((4, 3, 2), (4, 4, 4), 3)
Called OperationMemory.add_branch((4, 3, 2), (4, 4, 4), 3)
Called PredictionMemory.add_branch((4, 3, 3), (4, 4, 5), 3)
Called OperationMemory.add_branch((4, 3, 3), (4, 4, 5), 3)
Called PredictionMemory.add_branch((4, 3, 4), (4, 4, 6), 3)
Called OperationMemory.add_branch((4, 3, 4), (4, 4, 6), 3)

Do branching again, but with a different branching group. The first object should not produce a new branch at is already associated with ‘bg1’, but the second object should produce a new branch as it is not associated with any branching group yet:

>>> br._active_tasks[0,1,0,3] = True
>>> br._active_tasks[4,0,1,3] = True
>>> task = Mock('Task')
>>> task.branching_group = 'bg2'
>>> task.name = 'task3'
>>> data = Mock('Data')

Make sure that dummy Handler’s get_tind method returns target index with two objects (iteration, branch, object index, optional data start index, optional data stop index):

>>> data.get_tind.mock_returns = (
...     np.array([[0,1,0,0,0],
...               [4,0,1,0,0]], dtype=int), None)

Dummy Handler’s add_branch should return the target index for the new branches:

>>> data.add_branch.mock_returns = \
...     np.array([[4,5,1,0,0]], dtype=int)
>>> oper_mem = Mock('OperationMemory')
>>> pred_mem = Mock('PredictionMemory')
>>> br._no_op_branch = True
>>> data.linkage.links = {
...     (0,1): {(1,0): {0: [0], 2: [0], 3: [0,1,2,3]}},
...     (4,0): {(1,1): {0: [0], 2: [1], 3: [1,2,3,4]}},
...     (4,5): {(1,1): {0: [0], 2: [0], 3: [0,1,2,3]}}
... }

Start the actual branching:

>>> br.start(task, 'branch4', data, lexicon, 1, 0, pred_mem, oper_mem)
Called Data.get_tind(1, 0)
Called Data.block(1, 0, array([[0, 1, 0, 0, 0]]))
Called Data.add_branch(array([[4, 0, 1, 0, 0]]))
Called Data.block(1, 0, array([[4, 0, 1, 0, 0]]))
Called PredictionMemory.add_branch(array([4, 0, 1]), array([4, 5, 1]), 1)
Called OperationMemory.add_branch(array([4, 0, 1]), array([4, 5, 1]), 1)
Called PredictionMemory.add_branch((4, 0, 1), (4, 5, 0), 2)
Called OperationMemory.add_branch((4, 0, 1), (4, 5, 0), 2)
Called PredictionMemory.add_branch((4, 0, 1), (4, 5, 0), 3)
Called OperationMemory.add_branch((4, 0, 1), (4, 5, 0), 3)
Called PredictionMemory.add_branch((4, 0, 2), (4, 5, 1), 3)
Called OperationMemory.add_branch((4, 0, 2), (4, 5, 1), 3)
Called PredictionMemory.add_branch((4, 0, 3), (4, 5, 2), 3)
Called OperationMemory.add_branch((4, 0, 3), (4, 5, 2), 3)
Called PredictionMemory.add_branch((4, 0, 4), (4, 5, 3), 3)
Called OperationMemory.add_branch((4, 0, 4), (4, 5, 3), 3)
>>> br._parent_branches
array([[4, 0, 1, 0, 0]])
>>> br._new_branches
array([[4, 5, 1, 0, 0]])