'''Utility procedures for manipulating HANDE data.'''

import numpy as np

def groupby_beta_loops(data):
    '''Group a HANDE DMQMC data table by beta loop.

Parameters
----------
data : :class:`pandas.DataFrame`
    DMQMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.

Returns
-------
grouped : :class:`pandas.DataFrameGroupBy`
    GroupBy object with data table grouped by beta loop.
'''

    # Exploit the fact that (except for possibly the last beta loop due to wall
    # time) each beta loop contains the same set of iterations.
    indx = np.arange(len(data)) // len(data['iterations']).unique()
    return data.groupby(indx)

def groupby_iterations(data):
    '''Group a HANDE QMC data table by blocks of iterations.

Parameters
----------
data : :class:`pandas.DataFrame`
    QMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.

Returns
-------
grouped : :class:`pandas.DataFrameGroupBy`
    GroupBy object with data table grouped into blocks within which the
    iteration count increases monotonically.
'''
    indx = np.zeros(len(data))
    prev_iteration = -1
    curr_indx = 0
    for i in range(len(data)):
        if data['iterations'].iloc[i] < prev_iteration:
            # new block of iterations
            curr_indx += 1
        indx[i] = curr_indx
        prev_iteration = data['iterations'].iloc[i]

    return data.groupby(indx)

'''Utility procedures for manipulating HANDE data.'''

import numpy as np

def groupby_beta_loops(data, name='iterations'):
    '''Group a HANDE DMQMC data table by beta loop.

Parameters
----------
data : :class:`pandas.DataFrame`
    DMQMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.

Returns
-------
grouped : :class:`pandas.DataFrameGroupBy`
    GroupBy object with data table grouped by beta loop.
'''

    # Exploit the fact that (except for possibly the last beta loop due to wall
    # time) each beta loop contains the same set of iterations.
    indx = np.arange(len(data)) // len(data[name].unique())
    return data.groupby(indx)

def groupby_iterations(data):
    '''Group a HANDE QMC data table by blocks of iterations.

Parameters
----------
data : :class:`pandas.DataFrame`
    QMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.

Returns
-------
grouped : :class:`pandas.DataFrameGroupBy`
    GroupBy object with data table grouped into blocks within which the
    iteration count increases monotonically.
'''
    indx = np.zeros(len(data))
    prev_iteration = -1
    curr_indx = 0
    for i in range(len(data)):
        if data['iterations'].iloc[i] < prev_iteration:
            # new block of iterations
            curr_indx += 1
        indx[i] = curr_indx
        prev_iteration = data['iterations'].iloc[i]

    return data.groupby(indx)

Correct location of bracket so that grouping by beta loops is done correctly.
Correct location of bracket so that grouping by beta loops is done correctly.

  '''Utility procedures for manipulating HANDE data.'''
  
  import numpy as np
  
- def groupby_beta_loops(data):
+ def groupby_beta_loops(data, name='iterations'):
      '''Group a HANDE DMQMC data table by beta loop.
  
  Parameters
  ----------
  data : :class:`pandas.DataFrame`
      DMQMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.
  
  Returns
  -------
  grouped : :class:`pandas.DataFrameGroupBy`
      GroupBy object with data table grouped by beta loop.
  '''
  
      # Exploit the fact that (except for possibly the last beta loop due to wall
      # time) each beta loop contains the same set of iterations.
-     indx = np.arange(len(data)) // len(data['iterations']).unique()
+     indx = np.arange(len(data)) // len(data[name].unique())
      return data.groupby(indx)
  
  def groupby_iterations(data):
      '''Group a HANDE QMC data table by blocks of iterations.
  
  Parameters
  ----------
  data : :class:`pandas.DataFrame`
      QMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.
  
  Returns
  -------
  grouped : :class:`pandas.DataFrameGroupBy`
      GroupBy object with data table grouped into blocks within which the
      iteration count increases monotonically.
  '''
      indx = np.zeros(len(data))
      prev_iteration = -1
      curr_indx = 0
      for i in range(len(data)):
          if data['iterations'].iloc[i] < prev_iteration:
              # new block of iterations
              curr_indx += 1
          indx[i] = curr_indx
          prev_iteration = data['iterations'].iloc[i]
  
      return data.groupby(indx)
  
Correct location of bracket so that grouping by beta loops is done correctly.
## Code Before:
'''Utility procedures for manipulating HANDE data.'''

import numpy as np

def groupby_beta_loops(data):
    '''Group a HANDE DMQMC data table by beta loop.

Parameters
----------
data : :class:`pandas.DataFrame`
    DMQMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.

Returns
-------
grouped : :class:`pandas.DataFrameGroupBy`
    GroupBy object with data table grouped by beta loop.
'''

    # Exploit the fact that (except for possibly the last beta loop due to wall
    # time) each beta loop contains the same set of iterations.
    indx = np.arange(len(data)) // len(data['iterations']).unique()
    return data.groupby(indx)

def groupby_iterations(data):
    '''Group a HANDE QMC data table by blocks of iterations.

Parameters
----------
data : :class:`pandas.DataFrame`
    QMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.

Returns
-------
grouped : :class:`pandas.DataFrameGroupBy`
    GroupBy object with data table grouped into blocks within which the
    iteration count increases monotonically.
'''
    indx = np.zeros(len(data))
    prev_iteration = -1
    curr_indx = 0
    for i in range(len(data)):
        if data['iterations'].iloc[i] < prev_iteration:
            # new block of iterations
            curr_indx += 1
        indx[i] = curr_indx
        prev_iteration = data['iterations'].iloc[i]

    return data.groupby(indx)

## Instruction:
Correct location of bracket so that grouping by beta loops is done correctly.
## Code After:
'''Utility procedures for manipulating HANDE data.'''

import numpy as np

def groupby_beta_loops(data, name='iterations'):
    '''Group a HANDE DMQMC data table by beta loop.

Parameters
----------
data : :class:`pandas.DataFrame`
    DMQMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.

Returns
-------
grouped : :class:`pandas.DataFrameGroupBy`
    GroupBy object with data table grouped by beta loop.
'''

    # Exploit the fact that (except for possibly the last beta loop due to wall
    # time) each beta loop contains the same set of iterations.
    indx = np.arange(len(data)) // len(data[name].unique())
    return data.groupby(indx)

def groupby_iterations(data):
    '''Group a HANDE QMC data table by blocks of iterations.

Parameters
----------
data : :class:`pandas.DataFrame`
    QMC data table (e.g. obtained by :func:`pyhande.extract.extract_data`.

Returns
-------
grouped : :class:`pandas.DataFrameGroupBy`
    GroupBy object with data table grouped into blocks within which the
    iteration count increases monotonically.
'''
    indx = np.zeros(len(data))
    prev_iteration = -1
    curr_indx = 0
    for i in range(len(data)):
        if data['iterations'].iloc[i] < prev_iteration:
            # new block of iterations
            curr_indx += 1
        indx[i] = curr_indx
        prev_iteration = data['iterations'].iloc[i]

    return data.groupby(indx)
