nodes: 0, module; 1, function_definition; 2, function_name:basis; 3, parameters; 4, block; 5, identifier:start; 6, default_parameter; 7, default_parameter; 8, default_parameter; 9, default_parameter; 10, expression_statement; 11, if_statement; 12, expression_statement; 13, expression_statement; 14, expression_statement; 15, expression_statement; 16, if_statement; 17, if_statement; 18, expression_statement; 19, expression_statement; 20, expression_statement; 21, expression_statement; 22, for_statement; 23, expression_statement; 24, expression_statement; 25, expression_statement; 26, for_statement; 27, return_statement; 28, identifier:stop; 29, None; 30, identifier:dim; 31, integer:1; 32, identifier:sort; 33, string:"G"; 34, identifier:cross_truncation; 35, float:1.; 36, comment:"""
      Create an N-dimensional unit polynomial basis.
  
      Args:
          start (int, numpy.ndarray):
              the minimum polynomial to include. If int is provided, set as
              lowest total order.  If array of int, set as lower order along each
              axis.
          stop (int, numpy.ndarray):
              the maximum shape included. If omitted:
              ``stop <- start; start <- 0`` If int is provided, set as largest
              total order. If array of int, set as largest order along each axis.
          dim (int):
              dim of the basis. Ignored if array is provided in either start or
              stop.
          sort (str):
              The polynomial ordering where the letters ``G``, ``I`` and ``R``
              can be used to set grade, inverse and reverse to the ordering.  For
              ``basis(start=0, stop=2, dim=2, order=order)`` we get:
              ======  ==================
              order   output
              ======  ==================
              ""      [1 y y^2 x xy x^2]
              "G"     [1 y x y^2 xy x^2]
              "I"     [x^2 xy x y^2 y 1]
              "R"     [1 x x^2 y xy y^2]
              "GIR"   [y^2 xy x^2 y x 1]
              ======  ==================
          cross_truncation (float):
              Use hyperbolic cross truncation scheme to reduce the number of
              terms in expansion.
  
      Returns:
          (Poly) : Polynomial array.
  
      Examples:
          >>> print(chaospy.basis(4, 4, 2))
          [q0^4, q0^3q1, q0^2q1^2, q0q1^3, q1^4]
          >>> print(chaospy.basis([1, 1], [2, 2]))
          [q0q1, q0^2q1, q0q1^2, q0^2q1^2]
      """; 37, comparison_operator:stop is None; 38, block; 39, assignment; 40, assignment; 41, assignment; 42, assignment; 43, comparison_operator:start.size == 1; 44, block; 45, else_clause; 46, comparison_operator:stop.size == 1; 47, block; 48, else_clause; 49, assignment; 50, assignment; 51, assignment; 52, assignment; 53, identifier:idx; 54, identifier:pool; 55, block; 56, assignment; 57, assignment; 58, assignment; 59, identifier:I; 60, identifier:pool; 61, block; 62, call; 63, identifier:stop; 64, None; 65, expression_statement; 66, identifier:start; 67, call; 68, identifier:stop; 69, call; 70, identifier:dim; 71, call; 72, identifier:indices; 73, call; 74, attribute; 75, integer:1; 76, expression_statement; 77, block; 78, attribute; 79, integer:1; 80, expression_statement; 81, block; 82, identifier:pool; 83, call; 84, identifier:arg; 85, call; 86, subscript; 87, integer:1; 88, identifier:poly; 89, dictionary; 90, expression_statement; 91, expression_statement; 92, expression_statement; 93, identifier:x; 94, call; 95, subscript; 96, integer:1; 97, identifier:A; 98, dictionary; 99, expression_statement; 100, expression_statement; 101, expression_statement; 102, identifier:Poly; 103, argument_list; 104, assignment; 105, attribute; 106, argument_list; 107, attribute; 108, argument_list; 109, identifier:max; 110, argument_list; 111, attribute; 112, argument_list; 113, identifier:start; 114, identifier:size; 115, assignment; 116, expression_statement; 117, expression_statement; 118, identifier:stop; 119, identifier:size; 120, assignment; 121, expression_statement; 122, expression_statement; 123, identifier:list; 124, argument_list; 125, attribute; 126, argument_list; 127, identifier:arg; 128, integer:0; 129, assignment; 130, assignment; 131, assignment; 132, attribute; 133, argument_list; 134, identifier:x; 135, integer:0; 136, assignment; 137, assignment; 138, assignment; 139, identifier:A; 140, identifier:dim; 141, pattern_list; 142, expression_list; 143, identifier:numpy; 144, identifier:array; 145, identifier:start; 146, keyword_argument; 147, identifier:numpy; 148, identifier:array; 149, identifier:stop; 150, keyword_argument; 151, attribute; 152, attribute; 153, identifier:dim; 154, identifier:numpy; 155, identifier:array; 156, call; 157, identifier:bellow; 158, comparison_operator:numpy.sum(indices, -1) >= start; 159, assignment; 160, assignment; 161, identifier:above; 162, comparison_operator:numpy.sum(indices, -1) <= stop.item(); 163, assignment; 164, assignment; 165, subscript; 166, identifier:numpy; 167, identifier:zeros; 168, call; 169, keyword_argument; 170, identifier:idx; 171, call; 172, subscript; 173, identifier:arg; 174, identifier:arg; 175, call; 176, identifier:numpy; 177, identifier:zeros; 178, call; 179, keyword_argument; 180, identifier:I; 181, call; 182, subscript; 183, identifier:x; 184, identifier:x; 185, call; 186, identifier:start; 187, identifier:stop; 188, call; 189, identifier:start; 190, identifier:dtype; 191, identifier:int; 192, identifier:dtype; 193, identifier:int; 194, identifier:start; 195, identifier:size; 196, identifier:stop; 197, identifier:size; 198, attribute; 199, argument_list; 200, call; 201, identifier:start; 202, identifier:start; 203, binary_operator:numpy.ones(dim, dtype=int)*start; 204, identifier:bellow; 205, call; 206, call; 207, call; 208, identifier:stop; 209, binary_operator:numpy.ones(dim, dtype=int)*stop; 210, identifier:above; 211, call; 212, identifier:indices; 213, binary_operator:above*bellow; 214, identifier:len; 215, argument_list; 216, identifier:dtype; 217, identifier:int; 218, identifier:tuple; 219, argument_list; 220, identifier:poly; 221, identifier:idx; 222, attribute; 223, argument_list; 224, identifier:len; 225, argument_list; 226, identifier:dtype; 227, identifier:int; 228, identifier:tuple; 229, argument_list; 230, identifier:A; 231, identifier:I; 232, attribute; 233, argument_list; 234, attribute; 235, argument_list; 236, attribute; 237, identifier:bindex; 238, call; 239, binary_operator:2*numpy.max(stop); 240, identifier:dim; 241, identifier:sort; 242, identifier:cross_truncation; 243, attribute; 244, argument_list; 245, call; 246, identifier:start; 247, attribute; 248, argument_list; 249, attribute; 250, argument_list; 251, attribute; 252, argument_list; 253, call; 254, identifier:stop; 255, attribute; 256, argument_list; 257, identifier:above; 258, identifier:bellow; 259, identifier:pool; 260, identifier:idx; 261, identifier:numpy; 262, identifier:roll; 263, identifier:arg; 264, integer:1; 265, identifier:pool; 266, identifier:I; 267, identifier:numpy; 268, identifier:roll; 269, identifier:x; 270, integer:1; 271, identifier:numpy; 272, identifier:array; 273, integer:0; 274, identifier:chaospy; 275, identifier:bertran; 276, attribute; 277, argument_list; 278, integer:2; 279, call; 280, identifier:numpy; 281, identifier:sum; 282, identifier:indices; 283, unary_operator; 284, attribute; 285, argument_list; 286, identifier:numpy; 287, identifier:all; 288, comparison_operator:indices-start >= 0; 289, unary_operator; 290, identifier:numpy; 291, identifier:sum; 292, identifier:indices; 293, unary_operator; 294, identifier:stop; 295, identifier:item; 296, attribute; 297, argument_list; 298, identifier:numpy; 299, identifier:all; 300, comparison_operator:stop-indices >= 0; 301, unary_operator; 302, identifier:numpy; 303, identifier:min; 304, identifier:start; 305, attribute; 306, argument_list; 307, integer:1; 308, identifier:numpy; 309, identifier:ones; 310, identifier:dim; 311, keyword_argument; 312, binary_operator:indices-start; 313, integer:0; 314, integer:1; 315, integer:1; 316, identifier:numpy; 317, identifier:ones; 318, identifier:dim; 319, keyword_argument; 320, binary_operator:stop-indices; 321, integer:0; 322, integer:1; 323, identifier:numpy; 324, identifier:max; 325, identifier:stop; 326, identifier:dtype; 327, identifier:int; 328, identifier:indices; 329, identifier:start; 330, identifier:dtype; 331, identifier:int; 332, identifier:stop; 333, identifier:indices
code: def basis(start, stop=None, dim=1, sort="G", cross_truncation=1.):
      """
      Create an N-dimensional unit polynomial basis.
  
      Args:
          start (int, numpy.ndarray):
              the minimum polynomial to include. If int is provided, set as
              lowest total order.  If array of int, set as lower order along each
              axis.
          stop (int, numpy.ndarray):
              the maximum shape included. If omitted:
              ``stop <- start; start <- 0`` If int is provided, set as largest
              total order. If array of int, set as largest order along each axis.
          dim (int):
              dim of the basis. Ignored if array is provided in either start or
              stop.
          sort (str):
              The polynomial ordering where the letters ``G``, ``I`` and ``R``
              can be used to set grade, inverse and reverse to the ordering.  For
              ``basis(start=0, stop=2, dim=2, order=order)`` we get:
              ======  ==================
              order   output
              ======  ==================
              ""      [1 y y^2 x xy x^2]
              "G"     [1 y x y^2 xy x^2]
              "I"     [x^2 xy x y^2 y 1]
              "R"     [1 x x^2 y xy y^2]
              "GIR"   [y^2 xy x^2 y x 1]
              ======  ==================
          cross_truncation (float):
              Use hyperbolic cross truncation scheme to reduce the number of
              terms in expansion.
  
      Returns:
          (Poly) : Polynomial array.
  
      Examples:
          >>> print(chaospy.basis(4, 4, 2))
          [q0^4, q0^3q1, q0^2q1^2, q0q1^3, q1^4]
          >>> print(chaospy.basis([1, 1], [2, 2]))
          [q0q1, q0^2q1, q0q1^2, q0^2q1^2]
      """
      if stop is None:
          start, stop = numpy.array(0), start
  
      start = numpy.array(start, dtype=int)
      stop = numpy.array(stop, dtype=int)
      dim = max(start.size, stop.size, dim)
      indices = numpy.array(chaospy.bertran.bindex(
          numpy.min(start), 2*numpy.max(stop), dim, sort, cross_truncation))
  
      if start.size == 1:
          bellow = numpy.sum(indices, -1) >= start
  
      else:
          start = numpy.ones(dim, dtype=int)*start
          bellow = numpy.all(indices-start >= 0, -1)
  
  
      if stop.size == 1:
          above = numpy.sum(indices, -1) <= stop.item()
      else:
          stop = numpy.ones(dim, dtype=int)*stop
          above = numpy.all(stop-indices >= 0, -1)
  
      pool = list(indices[above*bellow])
  
      arg = numpy.zeros(len(pool), dtype=int)
      arg[0] = 1
      poly = {}
      for idx in pool:
          idx = tuple(idx)
          poly[idx] = arg
          arg = numpy.roll(arg, 1)
      x = numpy.zeros(len(pool), dtype=int)
      x[0] = 1
      A = {}
      for I in pool:
          I = tuple(I)
          A[I] = x
          x = numpy.roll(x,1)
  
      return Poly(A, dim)