nodes: 0, module; 1, function_definition; 2, function_name:make_input_dataframe_by_entity; 3, parameters; 4, block; 5, identifier:tax_benefit_system; 6, identifier:nb_persons; 7, identifier:nb_groups; 8, expression_statement; 9, expression_statement; 10, expression_statement; 11, expression_statement; 12, expression_statement; 13, expression_statement; 14, expression_statement; 15, comment:#; 16, expression_statement; 17, expression_statement; 18, comment:# A legacy role is an index that every person within an entity has.; 19, comment:# For instance, the 'first_parent' has legacy role 0, the 'second_parent' 1, the first 'child' 2, the second 3, etc.; 20, expression_statement; 21, expression_statement; 22, for_statement; 23, for_statement; 24, return_statement; 25, comment:"""
          Generate a dictionnary of dataframes containing nb_persons persons spread in nb_groups groups.
  
          :param TaxBenefitSystem tax_benefit_system: the tax_benefit_system to use
          :param int nb_persons: the number of persons in the system
          :param int nb_groups: the number of collective entities in the system
  
          :returns: A dictionary whose keys are entities and values the corresponding data frames
  
          Example:
  
          >>> from openfisca_survey_manager.input_dataframe_generator import make_input_dataframe_by_entity
          >>> from openfisca_country_template import CountryTaxBenefitSystem
          >>> tbs = CountryTaxBenefitSystem()
          >>> input_dataframe_by_entity = make_input_dataframe_by_entity(tbs, 400, 100)
          >>> sorted(input_dataframe_by_entity['person'].columns.tolist())
          ['household_id', 'household_legacy_role', 'household_role', 'person_id']
          >>> sorted(input_dataframe_by_entity['household'].columns.tolist())
          []
      """; 26, assignment; 27, assignment; 28, assignment; 29, assignment; 30, assignment; 31, call; 32, assignment; 33, assignment; 34, assignment; 35, assignment; 36, identifier:id_person; 37, call; 38, block; 39, identifier:entity; 40, attribute; 41, block; 42, identifier:input_dataframe_by_entity; 43, identifier:input_dataframe_by_entity; 44, call; 45, identifier:person_entity; 46, subscript; 47, identifier:person_id; 48, call; 49, identifier:input_dataframe_by_entity; 50, call; 51, subscript; 52, call; 53, attribute; 54, argument_list; 55, identifier:adults; 56, binary_operator:[0] + sorted(random.sample(range(1, nb_persons), nb_groups - 1)); 57, identifier:members_entity_id; 58, call; 59, identifier:members_legacy_role; 60, call; 61, identifier:id_group; 62, unary_operator; 63, identifier:range; 64, argument_list; 65, if_statement; 66, expression_statement; 67, expression_statement; 68, identifier:tax_benefit_system; 69, identifier:entities; 70, if_statement; 71, expression_statement; 72, expression_statement; 73, expression_statement; 74, expression_statement; 75, expression_statement; 76, expression_statement; 77, expression_statement; 78, identifier:dict; 79, argument_list; 80, list_comprehension; 81, integer:0; 82, attribute; 83, argument_list; 84, identifier:dict; 85, argument_list; 86, identifier:input_dataframe_by_entity; 87, attribute; 88, attribute; 89, argument_list; 90, subscript; 91, identifier:set_index; 92, string; 93, list; 94, call; 95, attribute; 96, argument_list; 97, attribute; 98, argument_list; 99, integer:1; 100, identifier:nb_persons; 101, comparison_operator:id_person in adults; 102, block; 103, else_clause; 104, assignment; 105, assignment; 106, attribute; 107, block; 108, assignment; 109, assignment; 110, assignment; 111, assignment; 112, assignment; 113, assignment; 114, call; 115, identifier:entity; 116, for_in_clause; 117, if_clause; 118, identifier:np; 119, identifier:arange; 120, identifier:nb_persons; 121, identifier:person_entity; 122, identifier:key; 123, identifier:pd; 124, identifier:DataFrame; 125, dictionary; 126, identifier:input_dataframe_by_entity; 127, attribute; 128, string_content:person_id; 129, integer:0; 130, identifier:sorted; 131, argument_list; 132, identifier:np; 133, identifier:empty; 134, identifier:nb_persons; 135, keyword_argument; 136, identifier:np; 137, identifier:empty; 138, identifier:nb_persons; 139, keyword_argument; 140, identifier:id_person; 141, identifier:adults; 142, expression_statement; 143, expression_statement; 144, block; 145, subscript; 146, identifier:legacy_role; 147, subscript; 148, identifier:id_group; 149, identifier:entity; 150, identifier:is_person; 151, continue_statement; 152, identifier:key; 153, attribute; 154, identifier:person_dataframe; 155, subscript; 156, subscript; 157, identifier:members_entity_id; 158, subscript; 159, identifier:members_legacy_role; 160, subscript; 161, call; 162, subscript; 163, call; 164, attribute; 165, argument_list; 166, identifier:entity; 167, attribute; 168, attribute; 169, pair; 170, identifier:person_entity; 171, identifier:key; 172, call; 173, identifier:dtype; 174, identifier:int; 175, identifier:dtype; 176, identifier:int; 177, augmented_assignment; 178, assignment; 179, expression_statement; 180, identifier:members_legacy_role; 181, identifier:id_person; 182, identifier:members_entity_id; 183, identifier:id_person; 184, identifier:entity; 185, identifier:key; 186, identifier:input_dataframe_by_entity; 187, attribute; 188, identifier:person_dataframe; 189, binary_operator:key + '_id'; 190, identifier:person_dataframe; 191, binary_operator:key + '_legacy_role'; 192, identifier:person_dataframe; 193, binary_operator:key + '_role'; 194, attribute; 195, argument_list; 196, identifier:input_dataframe_by_entity; 197, identifier:key; 198, attribute; 199, argument_list; 200, subscript; 201, identifier:set_index; 202, binary_operator:key + '_id'; 203, keyword_argument; 204, identifier:tax_benefit_system; 205, identifier:entities; 206, identifier:entity; 207, identifier:is_person; 208, binary_operator:person_entity.key + '_id'; 209, identifier:person_id; 210, attribute; 211, argument_list; 212, identifier:id_group; 213, integer:1; 214, identifier:legacy_role; 215, integer:0; 216, assignment; 217, identifier:person_entity; 218, identifier:key; 219, identifier:key; 220, string; 221, identifier:key; 222, string; 223, identifier:key; 224, string; 225, identifier:np; 226, identifier:where; 227, comparison_operator:members_legacy_role == 0; 228, attribute; 229, attribute; 230, identifier:pd; 231, identifier:DataFrame; 232, dictionary; 233, identifier:input_dataframe_by_entity; 234, identifier:key; 235, identifier:key; 236, string; 237, identifier:inplace; 238, True; 239, attribute; 240, string; 241, identifier:random; 242, identifier:sample; 243, call; 244, binary_operator:nb_groups - 1; 245, identifier:legacy_role; 246, conditional_expression:2 if legacy_role == 0 else legacy_role + 1; 247, string_content:_id; 248, string_content:_legacy_role; 249, string_content:_role; 250, identifier:members_legacy_role; 251, integer:0; 252, subscript; 253, identifier:key; 254, subscript; 255, identifier:key; 256, pair; 257, string_content:_id; 258, identifier:person_entity; 259, identifier:key; 260, string_content:_id; 261, identifier:range; 262, argument_list; 263, identifier:nb_groups; 264, integer:1; 265, integer:2; 266, comparison_operator:legacy_role == 0; 267, binary_operator:legacy_role + 1; 268, attribute; 269, integer:0; 270, attribute; 271, unary_operator; 272, binary_operator:key + '_id'; 273, call; 274, integer:1; 275, identifier:nb_persons; 276, identifier:legacy_role; 277, integer:0; 278, identifier:legacy_role; 279, integer:1; 280, identifier:entity; 281, identifier:flattened_roles; 282, identifier:entity; 283, identifier:flattened_roles; 284, integer:1; 285, identifier:key; 286, string; 287, identifier:range; 288, argument_list; 289, string_content:_id; 290, identifier:nb_groups
code: def make_input_dataframe_by_entity(tax_benefit_system, nb_persons, nb_groups):
      """
          Generate a dictionnary of dataframes containing nb_persons persons spread in nb_groups groups.
  
          :param TaxBenefitSystem tax_benefit_system: the tax_benefit_system to use
          :param int nb_persons: the number of persons in the system
          :param int nb_groups: the number of collective entities in the system
  
          :returns: A dictionary whose keys are entities and values the corresponding data frames
  
          Example:
  
          >>> from openfisca_survey_manager.input_dataframe_generator import make_input_dataframe_by_entity
          >>> from openfisca_country_template import CountryTaxBenefitSystem
          >>> tbs = CountryTaxBenefitSystem()
          >>> input_dataframe_by_entity = make_input_dataframe_by_entity(tbs, 400, 100)
          >>> sorted(input_dataframe_by_entity['person'].columns.tolist())
          ['household_id', 'household_legacy_role', 'household_role', 'person_id']
          >>> sorted(input_dataframe_by_entity['household'].columns.tolist())
          []
      """
      input_dataframe_by_entity = dict()
      person_entity = [entity for entity in tax_benefit_system.entities if entity.is_person][0]
      person_id = np.arange(nb_persons)
      input_dataframe_by_entity = dict()
      input_dataframe_by_entity[person_entity.key] = pd.DataFrame({
          person_entity.key + '_id': person_id,
          })
      input_dataframe_by_entity[person_entity.key].set_index('person_id')
      #
      adults = [0] + sorted(random.sample(range(1, nb_persons), nb_groups - 1))
      members_entity_id = np.empty(nb_persons, dtype = int)
      # A legacy role is an index that every person within an entity has.
      # For instance, the 'first_parent' has legacy role 0, the 'second_parent' 1, the first 'child' 2, the second 3, etc.
      members_legacy_role = np.empty(nb_persons, dtype = int)
      id_group = -1
      for id_person in range(nb_persons):
          if id_person in adults:
              id_group += 1
              legacy_role = 0
          else:
              legacy_role = 2 if legacy_role == 0 else legacy_role + 1
          members_legacy_role[id_person] = legacy_role
          members_entity_id[id_person] = id_group
  
      for entity in tax_benefit_system.entities:
          if entity.is_person:
              continue
          key = entity.key
          person_dataframe = input_dataframe_by_entity[person_entity.key]
          person_dataframe[key + '_id'] = members_entity_id
          person_dataframe[key + '_legacy_role'] = members_legacy_role
          person_dataframe[key + '_role'] = np.where(
              members_legacy_role == 0, entity.flattened_roles[0].key, entity.flattened_roles[-1].key)
          input_dataframe_by_entity[key] = pd.DataFrame({
              key + '_id': range(nb_groups)
              })
          input_dataframe_by_entity[key].set_index(key + '_id', inplace = True)
  
      return input_dataframe_by_entity