nodes: 0, module; 1, function_definition; 2, function_name:stratify; 3, parameters; 4, block; 5, identifier:self; 6, expression_statement; 7, expression_statement; 8, expression_statement; 9, if_statement; 10, function_definition; 11, expression_statement; 12, expression_statement; 13, expression_statement; 14, comment:"""
  		Stratifies the sample based on propensity score.
  		
  		By default the sample is divided into five equal-sized bins.
  		The number of bins can be set by modifying the object
  		attribute named blocks. Alternatively, custom-sized bins can
  		be created by setting blocks equal to a sorted list of numbers
  		between 0 and 1 indicating the bin boundaries.
  
  		This method should only be executed after the propensity score
  		has been estimated.
  		"""; 15, assignment; 16, assignment; 17, call; 18, block; 19, else_clause; 20, function_name:subset; 21, parameters; 22, block; 23, assignment; 24, assignment; 25, assignment; 26, pattern_list; 27, expression_list; 28, identifier:pscore; 29, subscript; 30, identifier:isinstance; 31, argument_list; 32, expression_statement; 33, block; 34, identifier:p_low; 35, identifier:p_high; 36, return_statement; 37, identifier:subsets; 38, list_comprehension; 39, identifier:strata; 40, list_comprehension; 41, attribute; 42, call; 43, identifier:Y; 44, identifier:D; 45, identifier:X; 46, subscript; 47, subscript; 48, subscript; 49, attribute; 50, string; 51, attribute; 52, identifier:int; 53, assignment; 54, expression_statement; 55, comment:# make a copy; should be sorted; 56, expression_statement; 57, comment:# avoids always dropping 1st unit; 58, binary_operator:(p_low < pscore) & (pscore <= p_high); 59, call; 60, for_in_clause; 61, call; 62, for_in_clause; 63, identifier:self; 64, identifier:strata; 65, identifier:Strata; 66, argument_list; 67, attribute; 68, string; 69, attribute; 70, string; 71, attribute; 72, string; 73, identifier:self; 74, identifier:raw_data; 75, string_content:pscore; 76, identifier:self; 77, identifier:blocks; 78, identifier:blocks; 79, call; 80, assignment; 81, assignment; 82, parenthesized_expression; 83, parenthesized_expression; 84, identifier:subset; 85, argument_list; 86, identifier:ps; 87, call; 88, identifier:CausalModel; 89, argument_list; 90, identifier:s; 91, identifier:subsets; 92, identifier:strata; 93, identifier:subsets; 94, identifier:pscore; 95, identifier:self; 96, identifier:raw_data; 97, string_content:Y; 98, identifier:self; 99, identifier:raw_data; 100, string_content:D; 101, identifier:self; 102, identifier:raw_data; 103, string_content:X; 104, identifier:split_equal_bins; 105, argument_list; 106, identifier:blocks; 107, subscript; 108, subscript; 109, integer:0; 110, comparison_operator:p_low < pscore; 111, comparison_operator:pscore <= p_high; 112, list_splat; 113, identifier:zip; 114, argument_list; 115, subscript; 116, subscript; 117, subscript; 118, identifier:pscore; 119, attribute; 120, attribute; 121, slice; 122, identifier:blocks; 123, integer:0; 124, identifier:p_low; 125, identifier:pscore; 126, identifier:pscore; 127, identifier:p_high; 128, identifier:ps; 129, identifier:blocks; 130, subscript; 131, identifier:Y; 132, identifier:s; 133, identifier:D; 134, identifier:s; 135, identifier:X; 136, identifier:s; 137, identifier:self; 138, identifier:blocks; 139, identifier:self; 140, identifier:blocks; 141, identifier:blocks; 142, slice; 143, integer:1
code: def stratify(self):
  
  		"""
  		Stratifies the sample based on propensity score.
  		
  		By default the sample is divided into five equal-sized bins.
  		The number of bins can be set by modifying the object
  		attribute named blocks. Alternatively, custom-sized bins can
  		be created by setting blocks equal to a sorted list of numbers
  		between 0 and 1 indicating the bin boundaries.
  
  		This method should only be executed after the propensity score
  		has been estimated.
  		"""
  
  		Y, D, X = self.raw_data['Y'], self.raw_data['D'], self.raw_data['X']
  		pscore = self.raw_data['pscore']
  
  		if isinstance(self.blocks, int):
  			blocks = split_equal_bins(pscore, self.blocks)
  		else:
  			blocks = self.blocks[:]  # make a copy; should be sorted
  			blocks[0] = 0  # avoids always dropping 1st unit
  
  		def subset(p_low, p_high):
  			return (p_low < pscore) & (pscore <= p_high)
  		subsets = [subset(*ps) for ps in zip(blocks, blocks[1:])]
  		strata = [CausalModel(Y[s], D[s], X[s]) for s in subsets]
  		self.strata = Strata(strata, subsets, pscore)