.. _brancher-py: ########### 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 # doctest: +ELLIPSIS 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)] # doctest: +ELLIPSIS (0, 'mock group', 'mock name', 'branch1', 'bg1', datetime.date(2002, 12, 31)) >>> br._branching_info[(4,3,0)] # doctest: +ELLIPSIS (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]])