parameters: (img, pad_size=16, mode=None, keep_batch_dim=False)
return_statement: return img3
docstring: pad parameterized 2d image based on the spherical positions of its vertices
      :param img: image to pad, with a shape [batch_size, H, W, ...] or [H, W] for a single image
      :param pad_size: size of the pad
      :param mode: lat or long or None, indicating if latitude dim proceeds longitude dim or the other way around
                   if mode is None, lat or long will be inferred automatically from the shape of the img,
                   with the assumption that lat dim < long dim
      :param keep_batch_dim: binary flag indicating whether to keep the batch dimension or not for single image
docstring_summary: pad parameterized 2d image based on the spherical positions of its vertices
      :param img: image to pad, with a shape [batch_size, H, W, ...] or [H, W] for a single image
      :param pad_size: size of the pad
      :param mode: lat or long or None, indicating if latitude dim proceeds longitude dim or the other way around
                   if mode is None, lat or long will be inferred automatically from the shape of the img,
                   with the assumption that lat dim < long dim
      :param keep_batch_dim: binary flag indicating whether to keep the batch dimension or not for single image
docstring_tokens: pad parameterized 2d image based on the spherical positions of its vertices : param img : image to pad with a shape [ batch_size H W ... ] or [ H W ] for a single image : param pad_size : size of the pad : param mode : lat or long or None indicating if latitude dim proceeds longitude dim or the other way around if mode is None lat or long will be inferred automatically from the shape of the img with the assumption that lat dim < long dim : param keep_batch_dim : binary flag indicating whether to keep the batch dimension or not for single image
function: def pad_2d_image_spherically(img, pad_size=16, mode=None, keep_batch_dim=False):
      """
      pad parameterized 2d image based on the spherical positions of its vertices
      :param img: image to pad, with a shape [batch_size, H, W, ...] or [H, W] for a single image
      :param pad_size: size of the pad
      :param mode: lat or long or None, indicating if latitude dim proceeds longitude dim or the other way around
                   if mode is None, lat or long will be inferred automatically from the shape of the img,
                   with the assumption that lat dim < long dim
      :param keep_batch_dim: binary flag indicating whether to keep the batch dimension or not for single image
      """
      if mode not in ['lat', 'long', None]:
          raise ValueError('mode should be either lat or long or None')
  
      # check if input is a single 2-D image or not, and pad to N-D if yes
      is_2d = False
      if img.ndim == 2:
          img = img[np.newaxis, ...]
          is_2d = True
  
      # set mode according to the input shape if it's None
      h, w = img.shape[1:3]
      if mode is None:
          if h < w:
              mode = 'lat'
          else:
              mode = 'long'
  
      # make sure lat dim is the 2nd dim and long dim is the 3rd dim
      dim_rest = list(np.arange(3, img.ndim))
      if mode == 'long':
          img = img.transpose([0, 2, 1, *dim_rest])
  
      if pad_size > 0:
          # pad the north pole on top
          top = img[:, 1:pad_size + 1, ...]  # get top pad without the first row (reflect)
          top = np.flip(top, axis=1)  # flip upside down
          top = np.roll(top, top.shape[2] // 2, axis=2)  # circularly shift by pi
  
          # similarly for the south pole on bottom
          bot = img[:, -pad_size - 1:-1, ...]
          bot = np.flip(bot, axis=1)
          bot = np.roll(bot, bot.shape[2] // 2, axis=2)
  
          # concatenate top and bottom before padding left and right
          img2 = np.concatenate((top, img, bot), axis=1)
  
          # pad left and right using the "wrap" function
          zero_pad_rest = ((0, 0),) * (img.ndim - 3)
          img3 = np.pad(img2, ((0, 0), (0, 0), (pad_size, pad_size), *zero_pad_rest), 'wrap')
      else:
          img3 = img
  
      # swap the dim back if it was swapped before
      if mode == 'long':
          img3 = img3.transpose([0, 2, 1, *dim_rest])
  
      # squeeze the tensor to 2-D if input is 2-D and the user does not want keep the extended dim
      if is_2d:
          if not keep_batch_dim:
              img3 = np.squeeze(img3)
  
      return img3
function_tokens: def pad_2d_image_spherically ( img , pad_size = 16 , mode = None , keep_batch_dim = False ) : if mode not in [ 'lat' , 'long' , None ] : raise ValueError ( 'mode should be either lat or long or None' ) # check if input is a single 2-D image or not, and pad to N-D if yes is_2d = False if img . ndim == 2 : img = img [ np . newaxis , ... ] is_2d = True # set mode according to the input shape if it's None h , w = img . shape [ 1 : 3 ] if mode is None : if h < w : mode = 'lat' else : mode = 'long' # make sure lat dim is the 2nd dim and long dim is the 3rd dim dim_rest = list ( np . arange ( 3 , img . ndim ) ) if mode == 'long' : img = img . transpose ( [ 0 , 2 , 1 , * dim_rest ] ) if pad_size > 0 : # pad the north pole on top top = img [ : , 1 : pad_size + 1 , ... ] # get top pad without the first row (reflect) top = np . flip ( top , axis = 1 ) # flip upside down top = np . roll ( top , top . shape [ 2 ] // 2 , axis = 2 ) # circularly shift by pi # similarly for the south pole on bottom bot = img [ : , - pad_size - 1 : - 1 , ... ] bot = np . flip ( bot , axis = 1 ) bot = np . roll ( bot , bot . shape [ 2 ] // 2 , axis = 2 ) # concatenate top and bottom before padding left and right img2 = np . concatenate ( ( top , img , bot ) , axis = 1 ) # pad left and right using the "wrap" function zero_pad_rest = ( ( 0 , 0 ) , ) * ( img . ndim - 3 ) img3 = np . pad ( img2 , ( ( 0 , 0 ) , ( 0 , 0 ) , ( pad_size , pad_size ) , * zero_pad_rest ) , 'wrap' ) else : img3 = img # swap the dim back if it was swapped before if mode == 'long' : img3 = img3 . transpose ( [ 0 , 2 , 1 , * dim_rest ] ) # squeeze the tensor to 2-D if input is 2-D and the user does not want keep the extended dim if is_2d : if not keep_batch_dim : img3 = np . squeeze ( img3 ) return img3