<commit_before># Copyright 2012 Daniel B. Allan
# dallan@pha.jhu.edu, daniel.b.allan@gmail.com
# http://pha.jhu.edu/~dallan
# http://www.danallan.com
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or (at
# your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses>.

from __future__ import division
import logging
import numpy as np
import pandas as pd
from pandas import DataFrame, Series
from scipy import interpolate
import pidly

logger = logging.getLogger(__name__)

def spline(t, pos, k=3, s=None):
    """Realize a Univariate spline, interpolating pos through all t.

    Parameters
    ----------
    t : Index (not Series!) of integers with possible gaps
    pos : DataFrame or Series of data to be interpolated
    k : integer
        polynomial order of spline, k <= 5
    s : float or tuple of floats (one for each column) or None
        Positive smoothing parameter used to determine the number of knots.
        For details, see documentation for scipy.interpolate.UnivariateSpline.

    Returns
    -------
    DataFrame of interpolated pos. The index is interpolated t.
    """
    first_frame, last_frame = t[[0, -1]]
    domain = np.arange(first_frame, 1 + last_frame)
    new_pos = []
    pos = DataFrame(pos) # in case pos is given as a Series
    import collections
    if not isinstance(s, collections.Iterable):
        s = [s]*len(pos.columns)
    assert len(s) == len(pos.columns), ("The length of s must match the number"
                                        "of columns in pos.")
    for i, col in enumerate(pos):
        spl = interpolate.UnivariateSpline(
            t.values, pos[col].values, k=k, s=s[i])
        new_col = Series(spl(domain), index=domain, name=col)
        new_pos.append(new_col)
    return pd.concat(new_pos, axis=1, keys=pos.columns)

def msd(traj, mpp, fps, max_lagtime=100, detail=False):
    """Compute the mean displacement and mean squared displacement of one
    trajectory over a range of time intervals.

    Parameters
    ----------
    traj : DataFrame of trajectories, including columns frame, x, and y
        If there is a probe column containing more than one value, an
        ensemble MSD will be computed. See imsd() to conveniently
        compute the MSD for each probe individually.
    mpp : microns per pixel
    fps : frames per second
    max_lagtime : intervals of frames out to which MSD is computed
        Default: 100
    detail : See below. Default False.

    Returns
    -------
    DataFrame([<x>, <y>, <x^2>, <y^2>, msd], index=t)

    If detail is True, the DataFrame also contains a column N,
    the estimated number of statistically independent measurements
    that comprise the result at each lagtime.

    Notes
    -----
    Input units are pixels and frames. Output units are microns and seconds.
    """
    pos = traj[['x', 'y']]
    t = traj['frame']
    # Reindex with consecutive frames, placing NaNs in the gaps.
    pos = pos.reindex(np.arange(pos.index[0], 1 + pos.index[-1]))
    max_lagtime = min(max_lagtime, len(t)) # checking to be safe
    lagtimes = 1 + np.arange(max_lagtime)
    disp = pd.concat([pos.sub(pos.shift(lt)) for lt in lagtimes],
                     keys=lagtimes, names=['lagt', 'frames'])
    results = mpp*disp.mean(level=0)
    results.columns = ['<x>', '<y>']
    results[['<x^2>', '<y^2>']] = mpp**2*(disp**2).mean(level=0)
    results['msd'] = mpp**2*(disp**2).mean(level=0).sum(1) # <r^2>
    # Estimated statistically independent measurements = 2N/t
    if detail:
        results['N'] = 2*disp.icol(0).count(level=0).div(Series(lagtimes))
    results['lagt'] = results.index/fps
    return results

def imsd(traj, mpp, fps, max_lagtime=100):
    """Compute the mean squared displacements of probes individually.

    Parameters
    ----------
    traj : DataFrame of trajectories of multiple probes, including
        columns probe, frame, x, and y
    mpp : microns per pixel
    fps : frames per second
    max_lagtime : intervals of frames out to which MSD is computed
        Default: 100

    Returns
    -------
    DataFrame([Probe 1 msd, Probe 2 msd, ...], index=t)

    Notes
    -----
    Input units are pixels and frames. Output units are microns and seconds.
    """
    ids = []
    msds = []
    for pid, ptraj in traj.groupby('probe'):
        msds.append(msd(ptraj, mpp, fps, max_lagtime, False))
        ids.append(pid)
    results = pd.concat(msds, keys=ids)
    # Swap MultiIndex levels so that unstack() makes probes into columns.
    results = results.swaplevel(0, 1)['msd'].unstack()
    return results

def emsd(traj, mpp, fps, max_lagtime=100, detail=False):
    """Compute the mean squared displacements of an ensemble of probes.

    Parameters
    ----------
    traj : DataFrame of trajectories of multiple probes, including
        columns probe, frame, x, and y
    mpp : microns per pixel
    fps : frames per second
    max_lagtime : intervals of frames out to which MSD is computed
        Default: 100

    Returns
    -------
    DataFrame([<x>, <y>, <x^2>, <y^2>, msd], index=t)

    Notes
    -----
    Input units are pixels and frames. Output units are microns and seconds.
    """
    ids = []
    msds = []
    for pid, ptraj in traj.groupby('probe'):
        msds.append(msd(ptraj, mpp, fps, max_lagtime, True))
        ids.append(pid)
    msds = pd.concat(msds, keys=ids, names=['probe', 'frame'])
    results = msds.mul(msds['N'], axis=0).mean(level=1) # weighted average
    results = results.div(msds['N'].mean(level=1), axis=0) # weights normalized
    # Above, lagt is lumped in with the rest for simplicity and speed.
    # Here, rebuild it from the frame index.
    results['lagt'] = results.index/fps
    del results['N']
    return results

def compute_drift(traj, smoothing=None):
    """Return the ensemble drift, x(t).

    Parameters
    ----------
    traj : DataFrame of trajectories, including columns x, y, frame, and probe
    smoothing : float, tuple of floats (x_smoothing, y_smoothing), or None
        Positive smoothing factor used to choose the number of knots.
        For details, see documentation for scipy.interpolate.UnivariateSpline.

    Returns
    -------
    drift : DataFrame([x, y], index=frame)

    Examples
    --------
    compute_drift(traj).plot() # Default smoothing usually smooths too much.
    compute_drift(traj, 0).plot() # not smoothed
    compute_drift(traj, 15).plot() # Try various smoothing values.
    compute_drift(traj, (18, 23)).plot() # Smooth x and y differently.

    drift = compute_drift(traj, (18, 23)) # Save good drift curves.
    corrected_traj = subtract_drift(traj, drift) # Apply them.
    """
    # Probe by probe, take the difference between frames.
    delta = pd.concat([t.set_index('frame', drop=False).diff()
                       for p, t in traj.groupby('probe')])
    # Keep only deltas between frames that are consecutive.
    delta = delta[delta['frame'] == 1]
    # Restore the original frame column (replacing delta frame).
    delta['frame'] = delta.index
    traj = delta.groupby('frame').mean()
    dx = spline(traj.index, traj[['x', 'y']], s=smoothing)
    x = dx.cumsum(0)
    return DataFrame(x, columns=['x', 'y'])

def subtract_drift(traj, drift=None):
    """Return a copy of probe trajectores with the overall drift subtracted out.

    Parameters
    ----------
    traj : DataFrame of trajectories, including columns x, y, and frame
    drift : optional DataFrame([x, y], index=frame) like output of
         compute_drift(). If no drift is passed, drift is computed from traj.

    Returns
    -------
    traj : a copy, having modified columns x and y
    """

    if drift is None:
        drift = compute_drift(traj)
    return traj.set_index('frame', drop=False).sub(drift, fill_value = 0)

def is_typical(msds, frame=23, lower=0.1, upper=0.9):
    """Examine individual probe MSDs, distinguishing outliers from those
    in the central quantile.

    Parameters
    ----------
    msds : DataFrame like the output of imsd()
        Columns correspond to probes, indexed by lagtime measured in frames.
    frame : integer frame number
        Compare MSDs at this lagtime. Default is 23 (1 second at 24 fps).
    lower : float between 0 and 1, default 0.1
        Probes with MSD up to this quantile are deemed outliers.
    upper : float between 0 and 1, default 0.9
        Probes with MSD above this quantile are deemed outliers.

    Returns
    -------
    Series of boolean values, indexed by probe number
    True = typical probe, False = outlier probe

    Example
    -------
    m = mr.imsd(traj, MPP, FPS)
    # Index by probe ID, slice using boolean output from is_typical(), and then
    # restore the original index, frame number.
    typical_traj = traj.set_index('probe').ix[is_typical(m)].reset_index()\
        .set_index('frame', drop=False)
    """
    a, b = msds.ix[frame].quantile(lower), msds.ix[frame].quantile(upper)
    return (msds.ix[frame] > a) & (msds.ix[frame] < b)

def idl_track(query, max_disp, min_appearances, memory=3):
    """Call Crocker/Weeks track.pro from IDL using pidly module.
    Returns one big array, where the last column is the probe ID."""
    idl = pidly.IDL()
    logger.info("Opened IDL process.")
    idl('pt = get_sql("{}")'.format(query))
    logger.info("IDL is done loading features from the database. Now tracking....")
    idl('t=track(pt, {}, goodenough={}, memory={})'.format(
        max_disp, min_appearances, memory))
    logger.info("IDL finished tracking. Now piping data into Python....")
    t = idl.ev('t')
    idl.close()
    return DataFrame(
        t, columns=['x', 'y', 'mass', 'size', 'ecc', 'frame', 'probe'])

def vanhove(msds, frame):
    """Compute the van Hove correlation function at given lagtime (frame span).

    Parameters
    ----------
    msds : DataFrame like the output of imsd()
        Columns correspond to probes, indexed by lagtime measured in frames.
    frame : integer frame number
        Compare MSDs at this lagtime. Default is 23 (1 second at 24 fps).

    Returns
    -------
    correlation : DataFrame with a correlation function for each probe in msds,
        indexed by displacement. That is, each column is the historgram, and
        the Index specifies the bins.
    """

def is_localized(traj, threshold=0.4):
    raise NotImplementedError, "I will rewrite this."

def is_diffusive(traj, threshold=0.9):
    raise NotImplementedError, "I will rewrite this."
<commit_msg>added statistic kw so imsd() can return <x><commit_after># Copyright 2012 Daniel B. Allan
# dallan@pha.jhu.edu, daniel.b.allan@gmail.com
# http://pha.jhu.edu/~dallan
# http://www.danallan.com
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or (at
# your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses>.

from __future__ import division
import logging
import numpy as np
import pandas as pd
from pandas import DataFrame, Series
from scipy import interpolate
import pidly

logger = logging.getLogger(__name__)

def spline(t, pos, k=3, s=None):
    """Realize a Univariate spline, interpolating pos through all t.

    Parameters
    ----------
    t : Index (not Series!) of integers with possible gaps
    pos : DataFrame or Series of data to be interpolated
    k : integer
        polynomial order of spline, k <= 5
    s : float or tuple of floats (one for each column) or None
        Positive smoothing parameter used to determine the number of knots.
        For details, see documentation for scipy.interpolate.UnivariateSpline.

    Returns
    -------
    DataFrame of interpolated pos. The index is interpolated t.
    """
    first_frame, last_frame = t[[0, -1]]
    domain = np.arange(first_frame, 1 + last_frame)
    new_pos = []
    pos = DataFrame(pos) # in case pos is given as a Series
    import collections
    if not isinstance(s, collections.Iterable):
        s = [s]*len(pos.columns)
    assert len(s) == len(pos.columns), ("The length of s must match the number"
                                        "of columns in pos.")
    for i, col in enumerate(pos):
        spl = interpolate.UnivariateSpline(
            t.values, pos[col].values, k=k, s=s[i])
        new_col = Series(spl(domain), index=domain, name=col)
        new_pos.append(new_col)
    return pd.concat(new_pos, axis=1, keys=pos.columns)

def msd(traj, mpp, fps, max_lagtime=100, detail=False):
    """Compute the mean displacement and mean squared displacement of one
    trajectory over a range of time intervals.

    Parameters
    ----------
    traj : DataFrame of trajectories, including columns frame, x, and y
        If there is a probe column containing more than one value, an
        ensemble MSD will be computed. See imsd() to conveniently
        compute the MSD for each probe individually.
    mpp : microns per pixel
    fps : frames per second
    max_lagtime : intervals of frames out to which MSD is computed
        Default: 100
    detail : See below. Default False.

    Returns
    -------
    DataFrame([<x>, <y>, <x^2>, <y^2>, msd], index=t)

    If detail is True, the DataFrame also contains a column N,
    the estimated number of statistically independent measurements
    that comprise the result at each lagtime.

    Notes
    -----
    Input units are pixels and frames. Output units are microns and seconds.
    """
    pos = traj[['x', 'y']]
    t = traj['frame']
    # Reindex with consecutive frames, placing NaNs in the gaps.
    pos = pos.reindex(np.arange(pos.index[0], 1 + pos.index[-1]))
    max_lagtime = min(max_lagtime, len(t)) # checking to be safe
    lagtimes = 1 + np.arange(max_lagtime)
    disp = pd.concat([pos.sub(pos.shift(lt)) for lt in lagtimes],
                     keys=lagtimes, names=['lagt', 'frames'])
    results = mpp*disp.mean(level=0)
    results.columns = ['<x>', '<y>']
    results[['<x^2>', '<y^2>']] = mpp**2*(disp**2).mean(level=0)
    results['msd'] = mpp**2*(disp**2).mean(level=0).sum(1) # <r^2>
    # Estimated statistically independent measurements = 2N/t
    if detail:
        results['N'] = 2*disp.icol(0).count(level=0).div(Series(lagtimes))
    results['lagt'] = results.index/fps
    return results

def imsd(traj, mpp, fps, max_lagtime=100, statistic='msd'):
    """Compute the mean squared displacements of probes individually.

    Parameters
    ----------
    traj : DataFrame of trajectories of multiple probes, including
        columns probe, frame, x, and y
    mpp : microns per pixel
    fps : frames per second
    max_lagtime : intervals of frames out to which MSD is computed
        Default: 100
    statistic : {'msd', '<x>', '<y>', '<x^2>', '<y^2>'}, default is 'msd'
        The functions msd() and emsd() return all these as columns. For
        imsd() you have to pick one.

    Returns
    -------
    DataFrame([Probe 1 msd, Probe 2 msd, ...], index=t)

    Notes
    -----
    Input units are pixels and frames. Output units are microns and seconds.
    """
    ids = []
    msds = []
    for pid, ptraj in traj.groupby('probe'):
        msds.append(msd(ptraj, mpp, fps, max_lagtime, False))
        ids.append(pid)
    results = pd.concat(msds, keys=ids)
    # Swap MultiIndex levels so that unstack() makes probes into columns.
    results = results.swaplevel(0, 1)[statistic].unstack()
    return results

def emsd(traj, mpp, fps, max_lagtime=100, detail=False):
    """Compute the mean squared displacements of an ensemble of probes.

    Parameters
    ----------
    traj : DataFrame of trajectories of multiple probes, including
        columns probe, frame, x, and y
    mpp : microns per pixel
    fps : frames per second
    max_lagtime : intervals of frames out to which MSD is computed
        Default: 100

    Returns
    -------
    DataFrame([<x>, <y>, <x^2>, <y^2>, msd], index=t)

    Notes
    -----
    Input units are pixels and frames. Output units are microns and seconds.
    """
    ids = []
    msds = []
    for pid, ptraj in traj.groupby('probe'):
        msds.append(msd(ptraj, mpp, fps, max_lagtime, True))
        ids.append(pid)
    msds = pd.concat(msds, keys=ids, names=['probe', 'frame'])
    results = msds.mul(msds['N'], axis=0).mean(level=1) # weighted average
    results = results.div(msds['N'].mean(level=1), axis=0) # weights normalized
    # Above, lagt is lumped in with the rest for simplicity and speed.
    # Here, rebuild it from the frame index.
    results['lagt'] = results.index/fps
    del results['N']
    return results

def compute_drift(traj, smoothing=None):
    """Return the ensemble drift, x(t).

    Parameters
    ----------
    traj : DataFrame of trajectories, including columns x, y, frame, and probe
    smoothing : float, tuple of floats (x_smoothing, y_smoothing), or None
        Positive smoothing factor used to choose the number of knots.
        For details, see documentation for scipy.interpolate.UnivariateSpline.

    Returns
    -------
    drift : DataFrame([x, y], index=frame)

    Examples
    --------
    compute_drift(traj).plot() # Default smoothing usually smooths too much.
    compute_drift(traj, 0).plot() # not smoothed
    compute_drift(traj, 15).plot() # Try various smoothing values.
    compute_drift(traj, (18, 23)).plot() # Smooth x and y differently.

    drift = compute_drift(traj, (18, 23)) # Save good drift curves.
    corrected_traj = subtract_drift(traj, drift) # Apply them.
    """
    # Probe by probe, take the difference between frames.
    delta = pd.concat([t.set_index('frame', drop=False).diff()
                       for p, t in traj.groupby('probe')])
    # Keep only deltas between frames that are consecutive.
    delta = delta[delta['frame'] == 1]
    # Restore the original frame column (replacing delta frame).
    delta['frame'] = delta.index
    traj = delta.groupby('frame').mean()
    dx = spline(traj.index, traj[['x', 'y']], s=smoothing)
    x = dx.cumsum(0)
    return DataFrame(x, columns=['x', 'y'])

def subtract_drift(traj, drift=None):
    """Return a copy of probe trajectores with the overall drift subtracted out.

    Parameters
    ----------
    traj : DataFrame of trajectories, including columns x, y, and frame
    drift : optional DataFrame([x, y], index=frame) like output of
         compute_drift(). If no drift is passed, drift is computed from traj.

    Returns
    -------
    traj : a copy, having modified columns x and y
    """

    if drift is None:
        drift = compute_drift(traj)
    return traj.set_index('frame', drop=False).sub(drift, fill_value = 0)

def is_typical(msds, frame=23, lower=0.1, upper=0.9):
    """Examine individual probe MSDs, distinguishing outliers from those
    in the central quantile.

    Parameters
    ----------
    msds : DataFrame like the output of imsd()
        Columns correspond to probes, indexed by lagtime measured in frames.
    frame : integer frame number
        Compare MSDs at this lagtime. Default is 23 (1 second at 24 fps).
    lower : float between 0 and 1, default 0.1
        Probes with MSD up to this quantile are deemed outliers.
    upper : float between 0 and 1, default 0.9
        Probes with MSD above this quantile are deemed outliers.

    Returns
    -------
    Series of boolean values, indexed by probe number
    True = typical probe, False = outlier probe

    Example
    -------
    m = mr.imsd(traj, MPP, FPS)
    # Index by probe ID, slice using boolean output from is_typical(), and then
    # restore the original index, frame number.
    typical_traj = traj.set_index('probe').ix[is_typical(m)].reset_index()\
        .set_index('frame', drop=False)
    """
    a, b = msds.ix[frame].quantile(lower), msds.ix[frame].quantile(upper)
    return (msds.ix[frame] > a) & (msds.ix[frame] < b)

def idl_track(query, max_disp, min_appearances, memory=3):
    """Call Crocker/Weeks track.pro from IDL using pidly module.
    Returns one big array, where the last column is the probe ID."""
    idl = pidly.IDL()
    logger.info("Opened IDL process.")
    idl('pt = get_sql("{}")'.format(query))
    logger.info("IDL is done loading features from the database. Now tracking....")
    idl('t=track(pt, {}, goodenough={}, memory={})'.format(
        max_disp, min_appearances, memory))
    logger.info("IDL finished tracking. Now piping data into Python....")
    t = idl.ev('t')
    idl.close()
    return DataFrame(
        t, columns=['x', 'y', 'mass', 'size', 'ecc', 'frame', 'probe'])

def vanhove(msds, frame):
    """Compute the van Hove correlation function at given lagtime (frame span).

    Parameters
    ----------
    msds : DataFrame like the output of imsd()
        Columns correspond to probes, indexed by lagtime measured in frames.
    frame : integer frame number
        Compare MSDs at this lagtime. Default is 23 (1 second at 24 fps).

    Returns
    -------
    correlation : DataFrame with a correlation function for each probe in msds,
        indexed by displacement. That is, each column is the historgram, and
        the Index specifies the bins.
    """

def is_localized(traj, threshold=0.4):
    raise NotImplementedError, "I will rewrite this."

def is_diffusive(traj, threshold=0.9):
    raise NotImplementedError, "I will rewrite this."
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import os
from osv import osv, fields
from mongodb_backend import osv_mongodb
from mongodb_backend.mongodb2 import mdbpool
from .erpwrapper import ErpWrapper
from yamlns import namespace as ns

from datetime import datetime
from plantmeter.resource import ProductionAggregator, ProductionPlant, ProductionMeter
from plantmeter.mongotimecurve import MongoTimeCurve, toLocal, asUtc
from plantmeter.isodates import isodate, naiveisodatetime, localisodate

class GenerationkwhProductionAggregator(osv.osv):
    """
    Serialization class for a plant mix. See plantmeter.aggregator.
    A mix would be a set of plants which production is to be added together.
    """

    _name = 'generationkwh.production.aggregator'

    _columns = {
        'name': fields.char('Name', size=50),
        'description': fields.char('Description', size=150),
        'enabled': fields.boolean('Enabled'),
        'plants': fields.one2many('generationkwh.production.plant', 'aggr_id', 'Plants')
    }

    _defaults = {
        'enabled': lambda *a: False
    }

    def get_kwh(self, cursor, uid, mix_id, start, end, context=None):
        '''Get production aggregation'''

        if not context:
            context = {}
        _aggr = self._createAggregator(cursor, uid, mix_id)
        return _aggr.get_kwh(start, end).tolist()

    def firstMeasurementDate(self, cursor, uid, mix_id, context=None):
        '''Get first measurement date'''

        if not context:
            context = {}
        _aggr = self._createAggregator(cursor, uid, mix_id)
        date = _aggr.firstMeasurementDate()
        return date if date else None

    def lastMeasurementDate(self, cursor, uid, mix_id, context=None):
        '''Get last measurement date'''

        if not context:
            context = {}
        _aggr = self._createAggregator(cursor, uid, mix_id)
        date = _aggr.lastMeasurementDate()
        return date if date else None

    def _createAggregator(self, cursor, uid, mix_id):
        def obj_to_dict(obj, attrs):
            return [(attr, getattr(obj, attr)) for attr in attrs]

        if isinstance(mix_id, list) or isinstance(mix_id, tuple):
            mix_id = mix_id[0]

        aggr = self.browse(cursor, uid, mix_id)
        curveProvider = MongoTimeCurve(mdbpool.get_db(),
            'tm_profile',
            creationField = 'create_date',
            timestampField = 'utc_gkwh_timestamp',
            )

        # TODO: Clean initialization method
        args = ['id', 'name', 'description', 'enabled']
        return ProductionAggregator(**dict(obj_to_dict(aggr, args) +
            dict(plants=[ProductionPlant(**dict(obj_to_dict(plant, args+['first_active_date']) +
                dict(meters=[ProductionMeter(
                    **dict(obj_to_dict(meter, args + ['uri','first_active_date']) +
                    dict(curveProvider=curveProvider).items()))
                for meter in plant.meters if meter.enabled]).items()))
            for plant in aggr.plants if plant.enabled]).items()))

GenerationkwhProductionAggregator()

class GenerationkwhProductionPlant(osv.osv):
    """
    Serialization class for a plant. See plantmeter.plant.
    """

    _name = 'generationkwh.production.plant'
    _columns = {
        'name': fields.char('Name', size=50),
        'description': fields.char('Description', size=150),
        'enabled': fields.boolean('Enabled'),
        'nshares': fields.integer('Number of shares'),
        'aggr_id': fields.many2one('generationkwh.production.aggregator', 'Production aggregator',
                                  required=True),
        'meters': fields.one2many('generationkwh.production.meter', 'plant_id', 'Meters'),
        'first_active_date': fields.date('First operative date'),
        'last_active_date': fields.date('Last operative date'),
    }
    _defaults = {
        'enabled': lambda *a: False,
    }

GenerationkwhProductionPlant()

class GenerationkwhProductionMeter(osv.osv):
    """
    Serialization class for a plant meter. See plantmeter.meter.
    """

    _name = 'generationkwh.production.meter'
    _columns = {
        'name': fields.char('Name', size=50),
        'description': fields.char('Description', size=150),
        'enabled': fields.boolean('Enabled'),
        'plant_id': fields.many2one('generationkwh.production.plant'),
        'uri': fields.char('Host', size=150, required=True),
        'first_active_date': fields.date('First operative date'),
        }
    _defaults = {
        'enabled': lambda *a: False,
    }

GenerationkwhProductionMeter()

class PlantShareProvider(ErpWrapper):
    """
    Provides a list for each plant of their active periods and their share value.
    To be used as provider for a LayeredShareCurve, that generates
    the curves to represent the share value of the active built plants.
    """
    def __init__(self, erp, cursor, uid, mixname, context=None):
        self.mixname = mixname
        super(PlantShareProvider, self).__init__(erp, cursor, uid, context)

    def items(self):
        Mix = self.erp.pool.get('generationkwh.production.aggregator')
        Plant = self.erp.pool.get('generationkwh.production.plant')
        mix_ids = Mix.search(self.cursor, self.uid, [
            ('name', '=', self.mixname)
        ])
        # if not mixids: ....

        plant_ids = Plant.search(self.cursor, self.uid, [
            ('aggr_id', '=', mix_ids[0]),
        ])
        plants = Plant.read(self.cursor, self.uid, plant_ids, [
        ])
        return [
            ns(
                mix=self.mixname,
                shares=plant['nshares'],
                firstEffectiveDate = isodate(plant['first_active_date']),
                lastEffectiveDate = isodate(plant['last_active_date']),
            )
            for plant in plants
        ]

class GenerationkwhProductionMeasurement(osv_mongodb.osv_mongodb):

    _name = 'generationkwh.production.measurement'
    _order = 'timestamp desc'

    _columns = {
        'name': fields.integer('Plant identifier'), # NOTE: workaround due mongodb backend
        'create_at': fields.datetime('Create datetime'),
        'datetime': fields.datetime('Exported datetime'),
        'daylight': fields.char('Exported datetime daylight',size=1),
        'ae': fields.float('Exported energy (kWh)')
    }

    def search(self, cursor, uid, args, offset=0, limit=0, order=None,
               context=None, count=False):
        '''Exact match for name.
        In mongodb, even when an index exists, not exact
        searches for fields scan all documents in collection
        making it extremely slow. Making name exact search
        reduces dramatically the amount of documents scanned'''

        new_args = []
        for arg in args:
            if not isinstance(arg, (list, tuple)):
                new_args.append(arg)
                continue
            field, operator, value = arg
            if field == 'name' and operator != '=':
                operator = '='
            new_args.append((field, operator, value))
        return super(GenerationkwhProductionMeasurement,
                     self).search(cursor, uid, new_args,
                                  offset=offset, limit=limit,
                                  order=order, context=context,
                                  count=count)

GenerationkwhProductionMeasurement()

class GenerationkwhProductionAggregatorTesthelper(osv.osv):
    '''Implements generationkwh production aggregation testhelper '''

    _name = 'generationkwh.production.aggregator.testhelper'
    _auto = False

    def get_kwh(self, cursor, uid, mix_id, start, end, context=None):
        mix = self.pool.get('generationkwh.production.aggregator')
        return mix.get_kwh(cursor, uid, mix_id,
                isodate(start), isodate(end), context)

    def firstMeasurementDate(self, cursor, uid, mix_id, context=None):
        mix = self.pool.get('generationkwh.production.aggregator')
        result = mix.firstMeasurementDate(cursor, uid, mix_id, context)
        return result and str(result)

    def lastMeasurementDate(self, cursor, uid, mix_id, context=None):
        production = self.pool.get('generationkwh.production.aggregator')
        result = production.lastMeasurementDate(cursor, uid, mix_id, context)
        return result and str(result)

    def clear_mongo_collections(self, cursor, uid, collections, context=None):
        for collection in collections:
            mdbpool.get_db().drop_collection(collection)

    def fillMeasurements(self, cursor, uid, first_date, meter_name, values):
        curveProvider = MongoTimeCurve(mdbpool.get_db(),
            'tm_profile',
            creationField = 'create_date',
            timestampField = 'utc_gkwh_timestamp',
            )
        curveProvider.update(
            start = localisodate(first_date),
            filter = meter_name,
            field = 'ae',
            data = values,
            )

    def fillMeasurementPoint(self, cursor, uid, pointTime, name, value, context=None):
        curveProvider = MongoTimeCurve(mdbpool.get_db(),
            'tm_profile',
            creationField = 'create_date',
            timestampField = 'utc_gkwh_timestamp',
            )
        curveProvider.fillPoint(
            datetime=toLocal(asUtc(naiveisodatetime(pointTime))),
            name=name,
            ae=value)

    def plantShareItems(self, cursor, uid, mixname):
        provider = PlantShareProvider(self, cursor, uid, mixname, context={})
        return [
            dict(
                item,
                lastEffectiveDate=str(item.lastEffectiveDate),
                firstEffectiveDate=str(item.firstEffectiveDate),
            )
            for item in provider.items()
        ]

GenerationkwhProductionAggregatorTesthelper()

# vim: ts=4 sw=4 et
<commit_msg>obj_to_dict -> attr_tuples<commit_after># -*- coding: utf-8 -*-

import os
from osv import osv, fields
from mongodb_backend import osv_mongodb
from mongodb_backend.mongodb2 import mdbpool
from .erpwrapper import ErpWrapper
from yamlns import namespace as ns

from datetime import datetime
from plantmeter.resource import ProductionAggregator, ProductionPlant, ProductionMeter
from plantmeter.mongotimecurve import MongoTimeCurve, toLocal, asUtc
from plantmeter.isodates import isodate, naiveisodatetime, localisodate

class GenerationkwhProductionAggregator(osv.osv):
    """
    Serialization class for a plant mix. See plantmeter.aggregator.
    A mix would be a set of plants which production is to be added together.
    """

    _name = 'generationkwh.production.aggregator'

    _columns = {
        'name': fields.char('Name', size=50),
        'description': fields.char('Description', size=150),
        'enabled': fields.boolean('Enabled'),
        'plants': fields.one2many('generationkwh.production.plant', 'aggr_id', 'Plants')
    }

    _defaults = {
        'enabled': lambda *a: False
    }

    def get_kwh(self, cursor, uid, mix_id, start, end, context=None):
        '''Get production aggregation'''

        if not context:
            context = {}
        _aggr = self._createAggregator(cursor, uid, mix_id)
        return _aggr.get_kwh(start, end).tolist()

    def firstMeasurementDate(self, cursor, uid, mix_id, context=None):
        '''Get first measurement date'''

        if not context:
            context = {}
        _aggr = self._createAggregator(cursor, uid, mix_id)
        date = _aggr.firstMeasurementDate()
        return date if date else None

    def lastMeasurementDate(self, cursor, uid, mix_id, context=None):
        '''Get last measurement date'''

        if not context:
            context = {}
        _aggr = self._createAggregator(cursor, uid, mix_id)
        date = _aggr.lastMeasurementDate()
        return date if date else None

    def _createAggregator(self, cursor, uid, mix_id):
        def attr_tuples(obj, attrs):
            # extracts name value tuples from erp browse object
            return [(attr, getattr(obj, attr)) for attr in attrs]

        if isinstance(mix_id, list) or isinstance(mix_id, tuple):
            mix_id = mix_id[0]

        aggr = self.browse(cursor, uid, mix_id)
        curveProvider = MongoTimeCurve(mdbpool.get_db(),
            'tm_profile',
            creationField = 'create_date',
            timestampField = 'utc_gkwh_timestamp',
            )

        # TODO: Clean initialization method
        args = ['id', 'name', 'description', 'enabled']
        return ProductionAggregator(**dict(attr_tuples(aggr, args) +
            dict(plants=[ProductionPlant(**dict(attr_tuples(plant, args+['first_active_date']) +
                dict(meters=[ProductionMeter(
                    **dict(attr_tuples(meter, args + ['uri','first_active_date']) +
                    dict(curveProvider=curveProvider).items()))
                for meter in plant.meters if meter.enabled]).items()))
            for plant in aggr.plants if plant.enabled]).items()))

GenerationkwhProductionAggregator()

class GenerationkwhProductionPlant(osv.osv):
    """
    Serialization class for a plant. See plantmeter.plant.
    """

    _name = 'generationkwh.production.plant'
    _columns = {
        'name': fields.char('Name', size=50),
        'description': fields.char('Description', size=150),
        'enabled': fields.boolean('Enabled'),
        'nshares': fields.integer('Number of shares'),
        'aggr_id': fields.many2one('generationkwh.production.aggregator', 'Production aggregator',
                                  required=True),
        'meters': fields.one2many('generationkwh.production.meter', 'plant_id', 'Meters'),
        'first_active_date': fields.date('First operative date'),
        'last_active_date': fields.date('Last operative date'),
    }
    _defaults = {
        'enabled': lambda *a: False,
    }

GenerationkwhProductionPlant()

class GenerationkwhProductionMeter(osv.osv):
    """
    Serialization class for a plant meter. See plantmeter.meter.
    """

    _name = 'generationkwh.production.meter'
    _columns = {
        'name': fields.char('Name', size=50),
        'description': fields.char('Description', size=150),
        'enabled': fields.boolean('Enabled'),
        'plant_id': fields.many2one('generationkwh.production.plant'),
        'uri': fields.char('Host', size=150, required=True),
        'first_active_date': fields.date('First operative date'),
        }
    _defaults = {
        'enabled': lambda *a: False,
    }

GenerationkwhProductionMeter()

class PlantShareProvider(ErpWrapper):
    """
    Provides a list for each plant of their active periods and their share value.
    To be used as provider for a LayeredShareCurve, that generates
    the curves to represent the share value of the active built plants.
    """
    def __init__(self, erp, cursor, uid, mixname, context=None):
        self.mixname = mixname
        super(PlantShareProvider, self).__init__(erp, cursor, uid, context)

    def items(self):
        Mix = self.erp.pool.get('generationkwh.production.aggregator')
        Plant = self.erp.pool.get('generationkwh.production.plant')
        mix_ids = Mix.search(self.cursor, self.uid, [
            ('name', '=', self.mixname)
        ])
        # if not mixids: ....

        plant_ids = Plant.search(self.cursor, self.uid, [
            ('aggr_id', '=', mix_ids[0]),
        ])
        plants = Plant.read(self.cursor, self.uid, plant_ids, [
        ])
        return [
            ns(
                mix=self.mixname,
                shares=plant['nshares'],
                firstEffectiveDate = isodate(plant['first_active_date']),
                lastEffectiveDate = isodate(plant['last_active_date']),
            )
            for plant in plants
        ]

class GenerationkwhProductionMeasurement(osv_mongodb.osv_mongodb):

    _name = 'generationkwh.production.measurement'
    _order = 'timestamp desc'

    _columns = {
        'name': fields.integer('Plant identifier'), # NOTE: workaround due mongodb backend
        'create_at': fields.datetime('Create datetime'),
        'datetime': fields.datetime('Exported datetime'),
        'daylight': fields.char('Exported datetime daylight',size=1),
        'ae': fields.float('Exported energy (kWh)')
    }

    def search(self, cursor, uid, args, offset=0, limit=0, order=None,
               context=None, count=False):
        '''Exact match for name.
        In mongodb, even when an index exists, not exact
        searches for fields scan all documents in collection
        making it extremely slow. Making name exact search
        reduces dramatically the amount of documents scanned'''

        new_args = []
        for arg in args:
            if not isinstance(arg, (list, tuple)):
                new_args.append(arg)
                continue
            field, operator, value = arg
            if field == 'name' and operator != '=':
                operator = '='
            new_args.append((field, operator, value))
        return super(GenerationkwhProductionMeasurement,
                     self).search(cursor, uid, new_args,
                                  offset=offset, limit=limit,
                                  order=order, context=context,
                                  count=count)

GenerationkwhProductionMeasurement()

class GenerationkwhProductionAggregatorTesthelper(osv.osv):
    '''Implements generationkwh production aggregation testhelper '''

    _name = 'generationkwh.production.aggregator.testhelper'
    _auto = False

    def get_kwh(self, cursor, uid, mix_id, start, end, context=None):
        mix = self.pool.get('generationkwh.production.aggregator')
        return mix.get_kwh(cursor, uid, mix_id,
                isodate(start), isodate(end), context)

    def firstMeasurementDate(self, cursor, uid, mix_id, context=None):
        mix = self.pool.get('generationkwh.production.aggregator')
        result = mix.firstMeasurementDate(cursor, uid, mix_id, context)
        return result and str(result)

    def lastMeasurementDate(self, cursor, uid, mix_id, context=None):
        production = self.pool.get('generationkwh.production.aggregator')
        result = production.lastMeasurementDate(cursor, uid, mix_id, context)
        return result and str(result)

    def clear_mongo_collections(self, cursor, uid, collections, context=None):
        for collection in collections:
            mdbpool.get_db().drop_collection(collection)

    def fillMeasurements(self, cursor, uid, first_date, meter_name, values):
        curveProvider = MongoTimeCurve(mdbpool.get_db(),
            'tm_profile',
            creationField = 'create_date',
            timestampField = 'utc_gkwh_timestamp',
            )
        curveProvider.update(
            start = localisodate(first_date),
            filter = meter_name,
            field = 'ae',
            data = values,
            )

    def fillMeasurementPoint(self, cursor, uid, pointTime, name, value, context=None):
        curveProvider = MongoTimeCurve(mdbpool.get_db(),
            'tm_profile',
            creationField = 'create_date',
            timestampField = 'utc_gkwh_timestamp',
            )
        curveProvider.fillPoint(
            datetime=toLocal(asUtc(naiveisodatetime(pointTime))),
            name=name,
            ae=value)

    def plantShareItems(self, cursor, uid, mixname):
        provider = PlantShareProvider(self, cursor, uid, mixname, context={})
        return [
            dict(
                item,
                lastEffectiveDate=str(item.lastEffectiveDate),
                firstEffectiveDate=str(item.firstEffectiveDate),
            )
            for item in provider.items()
        ]

GenerationkwhProductionAggregatorTesthelper()

# vim: ts=4 sw=4 et
<|endoftext|>


<commit_before># Mrs
# Copyright 2008-2012 Brigham Young University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""MapReduce Tasks.

A Task represents a unit of work and the mechanism for carrying it out.
"""

from __future__ import division, print_function

import copy
import itertools
from operator import itemgetter

from . import datasets
from . import fileformats
from . import serializers
from . import util

from logging import getLogger
logger = getLogger('mrs')

class Task(object):
    """Manage input and output for a piece of a map or reduce operation.

    The `task_index` determines the split of the input dataset that will be
    used by this Task, as well as the source number that will be created by
    this Task.
    """
    def __init__(self, op, input_ds, dataset_id, task_index, splits, storage,
            ext, serializers):
        self.op = op
        self.input_ds = input_ds
        self.dataset_id = dataset_id
        self.task_index = task_index
        self.splits = splits
        self.storage = storage if (storage is not None) else ''
        self.ext = ext
        self.serializers = serializers

        self.outdir = None
        self.output = None

    def outurls(self):
        return [(b.split, b.url) for b in self.output[:, :] if b.url]

    @staticmethod
    def from_op(op, *args):
        return op.task_class(op, *args)

    @staticmethod
    def from_args(op_args, urls, dataset_id, task_index, splits, storage,
            ext, input_ser_names, ser_names, program):
        """Converts from a simple tuple to a Task.

        The elements of the tuple correspond to the arguments of the
        Task.__init__ method, with the difference that the first argument
        is an Operation args tuple, and the second is a list of urls.
        """
        op = Operation.from_args(*op_args)

        input_serializers = serializers.from_names(input_ser_names, program)
        output_serializers = serializers.from_names(ser_names, program)

        input_ds = datasets.FileData(urls, program, splits=1,
                first_split=task_index, serializers=input_serializers)
        return Task.from_op(op, input_ds, dataset_id, task_index, splits,
                storage, ext, output_serializers)

    def to_args(self):
        """Converts the Task to a simple tuple.

        The elements of the tuple correspond to arguments of the Task.__init__
        method.  The first two elements of the tuples are lists of strings.
        The first is a list-of-strings representation of an operation, and the
        second is a list of urls.  The remaining elements are identical
        to the corresponding elements of the init method.
        """
        op_args = self.op.to_args()
        urls = [b.url for b in self.input_ds[:, self.task_index] if b.url]

        input_serializers = self.input_ds.serializers

        if input_serializers:
            input_ser_names = (input_serializers.key_s_name,
                    input_serializers.value_s_name)
        else:
            input_ser_names = ''

        if self.serializers:
            ser_names = (self.serializers.key_s_name,
                    self.serializers.value_s_name)
        else:
            ser_names = ''

        return (op_args, urls, self.dataset_id, self.task_index, self.splits,
                self.storage, self.ext, input_ser_names, ser_names)

    def _get_all_input(self, serial, sort=False, tmpdir=None,
            max_sort_size=None):
        """Returns an iterator over all input data."""
        if serial:
            self.input_ds.fetchall(_called_in_runner=True)
            uncopied_data = self.input_ds.data()
            # Avoid subtle race conditions in serial MapReduce when objects
            # are mutable by recklessly copying everything.  In the future, we
            # might make this disableable, but performance in serial MapReduce
            # is not a high priority.
            data = (copy.deepcopy(x) for x in uncopied_data)
            if sort:
                data = sorted(data, key=itemgetter(0))
        elif sort:
            sorted_ds = datasets.MergeSortData(self.input_ds, self.task_index,
                    max_sort_size, dir=tmpdir, _called_in_runner=True)
            data = sorted_ds.stream_data(_called_in_runner=True)
        else:
            data = self.input_ds.stream_split(self.task_index,
                    _called_in_runner=True)
        return data

    def _outdata_kwds(self, program, permanent, serial):
        """Returns arguments for the output dataset (common to all task types).
        """
        kwds = {'source': self.task_index,
                'parter': self.op.parter(program),
                'dir': self.outdir,
                'format': self.format(),
                'serializers': self.serializers,
                'splits': self.splits,
                }
        if not serial:
            kwds['write_only'] = True
        return kwds

    def make_outdir(self, default_dir):
        """Makes an output directory if necessary.

        Returns a bool indicating whether the files should be preserved (as
        opposed to automatically deleted).  Sets self.outdir, a path to a
        directory where output files should be created.
        """
        permanent = False
        if self.storage or default_dir:
            if self.storage:
                self.outdir = self.storage
                permanent = True
            else:
                self.outdir = util.mktempdir(default_dir, self.dataset_id + '_')

            if self.splits > 1:
                prefix = 'source_%s_' % self.task_index
                self.outdir = util.mktempdir(self.outdir, prefix)
        return permanent

    def format(self):
        """Return the write format given the Task's file extension."""
        if self.ext:
            format = fileformats.writerformat(self.ext)
        else:
            format = fileformats.default_write_format
        return format

class MapTask(Task):
    def run(self, program, default_dir, serial=False, max_sort_size=None):
        assert isinstance(self.op, MapOperation)

        all_input = self._get_all_input(serial)
        permanent = self.make_outdir(default_dir)
        kwds = self._outdata_kwds(program, permanent, serial)
        map_itr = self.op.map(program, all_input)
        self.output = datasets.LocalData(map_itr, permanent=permanent, **kwds)

class ReduceTask(Task):
    def run(self, program, default_dir, serial=False, max_sort_size=None):
        assert isinstance(self.op, ReduceOperation)

        all_input = self._get_all_input(serial, sort=True, tmpdir=default_dir,
                max_sort_size=max_sort_size)

        permanent = self.make_outdir(default_dir)
        kwds = self._outdata_kwds(program, permanent, serial)
        reduce_itr = self.op.reduce(program, all_input)
        self.output = datasets.LocalData(reduce_itr, permanent=permanent,
                **kwds)

class ReduceMapTask(Task):
    def run(self, program, default_dir, serial=False, max_sort_size=None):
        assert isinstance(self.op, ReduceMapOperation)

        all_input = self._get_all_input(serial, sort=True, tmpdir=default_dir,
                max_sort_size=max_sort_size)

        permanent = self.make_outdir(default_dir)
        kwds = self._outdata_kwds(program, permanent, serial)
        reduce_itr = self.op.reduce(program, all_input)
        map_itr = self.op.map(program, reduce_itr)
        self.output = datasets.LocalData(map_itr, permanent=permanent, **kwds)

class Operation(object):
    def __init__(self, part_name):
        self.part_name = part_name

    def parter(self, program):
        return getattr(program, self.part_name)

    @staticmethod
    def from_args(op_name, *args):
        cls = OP_CLASSES[op_name]
        return cls(*args)

class MapOperation(Operation):
    op_name = 'map'
    task_class = MapTask

    def __init__(self, map_name, combine_name, *args):
        Operation.__init__(self, *args)
        self.map_name = map_name
        self.combine_name = combine_name
        self.id = '%s' % self.map_name

    def map(self, program, input):
        """Yields map output iterating over the entries in input."""
        if self.map_name is None:
            mapper = None
        else:
            mapper = getattr(program, self.map_name)

        if self.combine_name:
            combine_op = ReduceOperation(self.combine_name, self.part_name)
        else:
            combine_op = None

        map_iter = self._map(mapper, input)
        if combine_op:
            # SORT PHASE
            sorted_map_iter = sorted(map_iter, key=itemgetter(0))
            combine_iter = combine_op.reduce(program, sorted_map_iter)
            return combine_iter
        else:
            return map_iter

    def _map(self, mapper, input):
        for inkey, invalue in input:
            for key, value in mapper(inkey, invalue):
                yield (key, value)

    def to_args(self):
        return (self.op_name, self.map_name, self.combine_name,
                self.part_name)

class ReduceOperation(Operation):
    op_name = 'reduce'
    task_class = ReduceTask

    def __init__(self, reduce_name, *args):
        Operation.__init__(self, *args)
        self.reduce_name = reduce_name
        self.id = '%s' % self.reduce_name

    def reduce(self, program, input):
        """Yields reduce output iterating over the entries in input.

        A reducer is an iterator taking a key and an iterator over values for
        that key.  It yields values for that key.
        """
        if self.reduce_name is None:
            reducer = None
        else:
            reducer = getattr(program, self.reduce_name)

        grouped_input = ((k, (pair[1] for pair in v)) for k, v in
            itertools.groupby(input, key=itemgetter(0)))

        for key, iterator in grouped_input:
            for value in reducer(key, iterator):
                yield (key, value)

    def to_args(self):
        return (self.op_name, self.reduce_name, self.part_name)

class ReduceMapOperation(MapOperation, ReduceOperation):
    op_name = 'reducemap'
    task_class = ReduceMapTask

    def __init__(self, reduce_name, map_name, combine_name, *args):
        Operation.__init__(self, *args)
        self.reduce_name = reduce_name
        self.map_name = map_name
        self.combine_name = combine_name
        self.id = '%s_%s' % (self.reduce_name, self.map_name)

    def to_args(self):
        return (self.op_name, self.reduce_name, self.map_name,
                self.combine_name, self.part_name)

OP_CLASSES = dict((op.op_name, op) for op in (MapOperation, ReduceOperation,
    ReduceMapOperation))

# vim: et sw=4 sts=4
<commit_msg>Make tasks delete merge sort temporary files when complete<commit_after># Mrs
# Copyright 2008-2012 Brigham Young University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""MapReduce Tasks.

A Task represents a unit of work and the mechanism for carrying it out.
"""

from __future__ import division, print_function

import copy
import itertools
from operator import itemgetter

from . import datasets
from . import fileformats
from . import serializers
from . import util

from logging import getLogger
logger = getLogger('mrs')

class Task(object):
    """Manage input and output for a piece of a map or reduce operation.

    The `task_index` determines the split of the input dataset that will be
    used by this Task, as well as the source number that will be created by
    this Task.
    """
    def __init__(self, op, input_ds, dataset_id, task_index, splits, storage,
            ext, serializers):
        self.op = op
        self.input_ds = input_ds
        self.dataset_id = dataset_id
        self.task_index = task_index
        self.splits = splits
        self.storage = storage if (storage is not None) else ''
        self.ext = ext
        self.serializers = serializers

        self.outdir = None
        self.output = None
        self.sorted_ds = None

    def outurls(self):
        return [(b.split, b.url) for b in self.output[:, :] if b.url]

    @staticmethod
    def from_op(op, *args):
        return op.task_class(op, *args)

    @staticmethod
    def from_args(op_args, urls, dataset_id, task_index, splits, storage,
            ext, input_ser_names, ser_names, program):
        """Converts from a simple tuple to a Task.

        The elements of the tuple correspond to the arguments of the
        Task.__init__ method, with the difference that the first argument
        is an Operation args tuple, and the second is a list of urls.
        """
        op = Operation.from_args(*op_args)

        input_serializers = serializers.from_names(input_ser_names, program)
        output_serializers = serializers.from_names(ser_names, program)

        input_ds = datasets.FileData(urls, program, splits=1,
                first_split=task_index, serializers=input_serializers)
        return Task.from_op(op, input_ds, dataset_id, task_index, splits,
                storage, ext, output_serializers)

    def to_args(self):
        """Converts the Task to a simple tuple.

        The elements of the tuple correspond to arguments of the Task.__init__
        method.  The first two elements of the tuples are lists of strings.
        The first is a list-of-strings representation of an operation, and the
        second is a list of urls.  The remaining elements are identical
        to the corresponding elements of the init method.
        """
        op_args = self.op.to_args()
        urls = [b.url for b in self.input_ds[:, self.task_index] if b.url]

        input_serializers = self.input_ds.serializers

        if input_serializers:
            input_ser_names = (input_serializers.key_s_name,
                    input_serializers.value_s_name)
        else:
            input_ser_names = ''

        if self.serializers:
            ser_names = (self.serializers.key_s_name,
                    self.serializers.value_s_name)
        else:
            ser_names = ''

        return (op_args, urls, self.dataset_id, self.task_index, self.splits,
                self.storage, self.ext, input_ser_names, ser_names)

    def _get_all_input(self, serial, sort=False, default_dir=None,
            max_sort_size=None):
        """Returns an iterator over all input data."""
        if serial:
            self.input_ds.fetchall(_called_in_runner=True)
            uncopied_data = self.input_ds.data()
            # Avoid subtle race conditions in serial MapReduce when objects
            # are mutable by recklessly copying everything.  In the future, we
            # might make this disableable, but performance in serial MapReduce
            # is not a high priority.
            data = (copy.deepcopy(x) for x in uncopied_data)
            if sort:
                data = sorted(data, key=itemgetter(0))
        elif sort:
            tmpdir = util.mktempdir(default_dir, 'merge_%s_' % self.dataset_id)
            sorted_ds = datasets.MergeSortData(self.input_ds, self.task_index,
                    max_sort_size, dir=tmpdir, _called_in_runner=True)
            data = sorted_ds.stream_data(_called_in_runner=True)
            self.sorted_ds = sorted_ds
        else:
            data = self.input_ds.stream_split(self.task_index,
                    _called_in_runner=True)
        return data

    def _outdata_kwds(self, program, permanent, serial):
        """Returns arguments for the output dataset (common to all task types).
        """
        kwds = {'source': self.task_index,
                'parter': self.op.parter(program),
                'dir': self.outdir,
                'format': self.format(),
                'serializers': self.serializers,
                'splits': self.splits,
                }
        if not serial:
            kwds['write_only'] = True
        return kwds

    def make_outdir(self, default_dir):
        """Makes an output directory if necessary.

        Returns a bool indicating whether the files should be preserved (as
        opposed to automatically deleted).  Sets self.outdir, a path to a
        directory where output files should be created.
        """
        permanent = False
        if self.storage or default_dir:
            if self.storage:
                self.outdir = self.storage
                permanent = True
            else:
                self.outdir = util.mktempdir(default_dir, self.dataset_id + '_')

            if self.splits > 1:
                prefix = 'source_%s_' % self.task_index
                self.outdir = util.mktempdir(self.outdir, prefix)
        return permanent

    def format(self):
        """Return the write format given the Task's file extension."""
        if self.ext:
            format = fileformats.writerformat(self.ext)
        else:
            format = fileformats.default_write_format
        return format

class MapTask(Task):
    def run(self, program, default_dir, serial=False, max_sort_size=None):
        assert isinstance(self.op, MapOperation)

        all_input = self._get_all_input(serial)
        permanent = self.make_outdir(default_dir)
        kwds = self._outdata_kwds(program, permanent, serial)
        map_itr = self.op.map(program, all_input)
        self.output = datasets.LocalData(map_itr, permanent=permanent, **kwds)

class ReduceTask(Task):
    def run(self, program, default_dir, serial=False, max_sort_size=None):
        assert isinstance(self.op, ReduceOperation)

        all_input = self._get_all_input(serial, sort=True,
                default_dir=default_dir, max_sort_size=max_sort_size)

        permanent = self.make_outdir(default_dir)
        kwds = self._outdata_kwds(program, permanent, serial)
        reduce_itr = self.op.reduce(program, all_input)
        self.output = datasets.LocalData(reduce_itr, permanent=permanent,
                **kwds)
        if self.sorted_ds is not None:
            self.sorted_ds.delete()

class ReduceMapTask(Task):
    def run(self, program, default_dir, serial=False, max_sort_size=None):
        assert isinstance(self.op, ReduceMapOperation)

        all_input = self._get_all_input(serial, sort=True,
                default_dir=default_dir, max_sort_size=max_sort_size)

        permanent = self.make_outdir(default_dir)
        kwds = self._outdata_kwds(program, permanent, serial)
        reduce_itr = self.op.reduce(program, all_input)
        map_itr = self.op.map(program, reduce_itr)
        self.output = datasets.LocalData(map_itr, permanent=permanent, **kwds)
        if self.sorted_ds is not None:
            self.sorted_ds.delete()

class Operation(object):
    def __init__(self, part_name):
        self.part_name = part_name

    def parter(self, program):
        return getattr(program, self.part_name)

    @staticmethod
    def from_args(op_name, *args):
        cls = OP_CLASSES[op_name]
        return cls(*args)

class MapOperation(Operation):
    op_name = 'map'
    task_class = MapTask

    def __init__(self, map_name, combine_name, *args):
        Operation.__init__(self, *args)
        self.map_name = map_name
        self.combine_name = combine_name
        self.id = '%s' % self.map_name

    def map(self, program, input):
        """Yields map output iterating over the entries in input."""
        if self.map_name is None:
            mapper = None
        else:
            mapper = getattr(program, self.map_name)

        if self.combine_name:
            combine_op = ReduceOperation(self.combine_name, self.part_name)
        else:
            combine_op = None

        map_iter = self._map(mapper, input)
        if combine_op:
            # SORT PHASE
            sorted_map_iter = sorted(map_iter, key=itemgetter(0))
            combine_iter = combine_op.reduce(program, sorted_map_iter)
            return combine_iter
        else:
            return map_iter

    def _map(self, mapper, input):
        for inkey, invalue in input:
            for key, value in mapper(inkey, invalue):
                yield (key, value)

    def to_args(self):
        return (self.op_name, self.map_name, self.combine_name,
                self.part_name)

class ReduceOperation(Operation):
    op_name = 'reduce'
    task_class = ReduceTask

    def __init__(self, reduce_name, *args):
        Operation.__init__(self, *args)
        self.reduce_name = reduce_name
        self.id = '%s' % self.reduce_name

    def reduce(self, program, input):
        """Yields reduce output iterating over the entries in input.

        A reducer is an iterator taking a key and an iterator over values for
        that key.  It yields values for that key.
        """
        if self.reduce_name is None:
            reducer = None
        else:
            reducer = getattr(program, self.reduce_name)

        grouped_input = ((k, (pair[1] for pair in v)) for k, v in
            itertools.groupby(input, key=itemgetter(0)))

        for key, iterator in grouped_input:
            for value in reducer(key, iterator):
                yield (key, value)

    def to_args(self):
        return (self.op_name, self.reduce_name, self.part_name)

class ReduceMapOperation(MapOperation, ReduceOperation):
    op_name = 'reducemap'
    task_class = ReduceMapTask

    def __init__(self, reduce_name, map_name, combine_name, *args):
        Operation.__init__(self, *args)
        self.reduce_name = reduce_name
        self.map_name = map_name
        self.combine_name = combine_name
        self.id = '%s_%s' % (self.reduce_name, self.map_name)

    def to_args(self):
        return (self.op_name, self.reduce_name, self.map_name,
                self.combine_name, self.part_name)

OP_CLASSES = dict((op.op_name, op) for op in (MapOperation, ReduceOperation,
    ReduceMapOperation))

# vim: et sw=4 sts=4
<|endoftext|>


<commit_before>"""
Evaluates the estimated results of the Segmentation dataset against the
ground truth (human annotated data).
"""
import jams
from joblib import Parallel, delayed
import logging
import mir_eval
import numpy as np
import os
import pandas as pd
import six
import sys

# Local stuff
import msaf
from msaf.exceptions import NoReferencesError
import msaf.input_output as io
from msaf import utils

def print_results(results):
    """Print all the results.

    Parameters
    ----------
    results: pd.DataFrame
        Dataframe with all the results
    """
    if len(results) == 0:
        return
    res = results.mean()
    logging.info(results["track_id"])
    logging.info("Results:\n%s" % res)

def compute_results(ann_inter, est_inter, ann_labels, est_labels, bins,
                    est_file, weight=0.58):
    """Compute the results using all the available evaluations.

    Parameters
    ----------
    ann_inter : np.array
        Annotated intervals in seconds.
    est_inter : np.array
        Estimated intervals in seconds.
    ann_labels : np.array
        Annotated labels.
    est_labels : np.array
        Estimated labels
    bins : int
        Number of bins for the information gain.
    est_file : str
        Path to the output file to store results.
    weight: float
        Weight the Precision and Recall values of the hit rate boundaries
        differently (<1 will weight Precision higher, >1 will weight Recall
        higher).
        The default parameter (0.58) is the one proposed in (Nieto et al. 2014)

    Return
    ------
    results : dict
        Contains the results of all the evaluations for the given file.
        Keys are the following:
            track_id: Name of the track
            HitRate_3F: F-measure of hit rate at 3 seconds
            HitRate_3P: Precision of hit rate at 3 seconds
            HitRate_3R: Recall of hit rate at 3 seconds
            HitRate_0.5F: F-measure of hit rate at 0.5 seconds
            HitRate_0.5P: Precision of hit rate at 0.5 seconds
            HitRate_0.5R: Recall of hit rate at 0.5 seconds
            HitRate_w3F: F-measure of hit rate at 3 seconds weighted
            HitRate_w0.5F: F-measure of hit rate at 0.5 seconds weighted
            HitRate_wt3F: F-measure of hit rate at 3 seconds weighted and
                          trimmed
            HitRate_wt0.5F: F-measure of hit rate at 0.5 seconds weighted
                            and trimmed
            HitRate_t3F: F-measure of hit rate at 3 seconds (trimmed)
            HitRate_t3P: Precision of hit rate at 3 seconds (trimmed)
            HitRate_t3F: Recall of hit rate at 3 seconds (trimmed)
            HitRate_t0.5F: F-measure of hit rate at 0.5 seconds (trimmed)
            HitRate_t0.5P: Precision of hit rate at 0.5 seconds (trimmed)
            HitRate_t0.5R: Recall of hit rate at 0.5 seconds (trimmed)
            DevA2E: Median deviation of annotation to estimation
            DevE2A: Median deviation of estimation to annotation
            D: Information gain
            PWF: F-measure of pair-wise frame clustering
            PWP: Precision of pair-wise frame clustering
            PWR: Recall of pair-wise frame clustering
            Sf: F-measure normalized entropy score
            So: Oversegmentation normalized entropy score
            Su: Undersegmentation normalized entropy score
    """
    res = {}

    # --Boundaries-- #
    # Hit Rate standard
    res["HitRate_3P"], res["HitRate_3R"], res["HitRate_3F"] = \
        mir_eval.segment.detection(ann_inter, est_inter, window=3, trim=False)
    res["HitRate_0.5P"], res["HitRate_0.5R"], res["HitRate_0.5F"] = \
        mir_eval.segment.detection(ann_inter, est_inter, window=.5, trim=False)

    # Hit rate trimmed
    res["HitRate_t3P"], res["HitRate_t3R"], res["HitRate_t3F"] = \
        mir_eval.segment.detection(ann_inter, est_inter, window=3, trim=True)
    res["HitRate_t0.5P"], res["HitRate_t0.5R"], res["HitRate_t0.5F"] = \
        mir_eval.segment.detection(ann_inter, est_inter, window=.5, trim=True)

    # Hit rate weighted
    res["HitRate_w3P"], _, _ = mir_eval.segment.detection(
        ann_inter, est_inter, window=3, trim=False, beta=weight)
    res["HitRate_w0.5P"], _, _ = mir_eval.segment.detection(
        ann_inter, est_inter, window=.5, trim=False, beta=weight)

    # Hit rate weighted and trimmed
    res["HitRate_wt3P"], _, _ = mir_eval.segment.detection(
        ann_inter, est_inter, window=3, trim=True, beta=weight)
    res["HitRate_wt0.5P"], _, _ = mir_eval.segment.detection(
        ann_inter, est_inter, window=.5, trim=True, beta=weight)

    # Information gain
    res["D"] = compute_information_gain(ann_inter, est_inter, est_file,
                                        bins=bins)

    # Median Deviations
    res["DevR2E"], res["DevE2R"] = mir_eval.segment.deviation(
        ann_inter, est_inter, trim=False)
    res["DevtR2E"], res["DevtE2R"] = mir_eval.segment.deviation(
        ann_inter, est_inter, trim=True)

    # --Labels-- #
    if est_labels is not None and "-1" in est_labels:
        est_labels = None

    if est_labels is not None and len(est_labels) != 0:
        try:
            # Align labels with intervals
            ann_labels = list(ann_labels)
            est_labels = list(est_labels)
            ann_inter, ann_labels = mir_eval.util.adjust_intervals(ann_inter,
                                                                   ann_labels)
            est_inter, est_labels = mir_eval.util.adjust_intervals(
                est_inter, est_labels, t_min=0, t_max=ann_inter.max())

            # Pair-wise frame clustering
            res["PWP"], res["PWR"], res["PWF"] = mir_eval.segment.pairwise(
                ann_inter, ann_labels, est_inter, est_labels)

            # Normalized Conditional Entropies
            res["So"], res["Su"], res["Sf"] = mir_eval.segment.nce(
                ann_inter, ann_labels, est_inter, est_labels)
        except:
            logging.warning("Labeling evaluation failed in file: %s" %
                            est_file)
            return {}

    # Names
    base = os.path.basename(est_file)
    res["track_id"] = base[:-5]

    return res

def compute_gt_results(est_file, ref_file, boundaries_id, labels_id, config,
                       bins=251, annotator_id=0):
    """Computes the results by using the ground truth dataset identified by
    the annotator parameter.

    Return
    ------
    results : dict
        Dictionary of the results (see function compute_results).
    """
    try:
        if config["hier"]:
            ref_times, ref_labels, ref_levels = \
                msaf.io.read_hier_references(
                    ref_file, annotation_id=0,
                    exclude_levels=["segment_salami_function"])
        else:
            jam = jams.load(ref_file, validate=False)
            ann = jam.search(namespace='segment_.*')[annotator_id]
            ref_inter, ref_labels = ann.data.to_interval_values()
    except:
        logging.warning("No references for file: %s" % ref_file)
        return {}

    # Read estimations with correct configuration
    est_inter, est_labels = io.read_estimations(est_file, boundaries_id,
                                                labels_id, **config)
    if len(est_inter) == 0:
        logging.warning("No estimations for file: %s" % est_file)
        return {}

    # Compute the results and return
    logging.info("Evaluating %s" % os.path.basename(est_file))
    if config["hier"]:
        # Hierarchical
        assert len(est_inter) == len(est_labels), "Same number of levels " \
            "are required in the boundaries and labels for the hierarchical " \
            "evaluation."
        est_times = []
        est_labels = []

        # Sort based on how many segments per level
        est_inter = sorted(est_inter, key=lambda level: len(level))

        for inter in est_inter:
            est_times.append(msaf.utils.intervals_to_times(inter))
            # Add fake labels (hierarchical eval does not use labels --yet--)
            est_labels.append(np.ones(len(est_times[-1]) - 1) * -1)

        # Align the times
        utils.align_end_hierarchies(est_times, ref_times, thres=1)

        # To intervals
        est_hier = [utils.times_to_intervals(times) for times in est_times]
        ref_hier = [utils.times_to_intervals(times) for times in ref_times]

        # Compute evaluations
        res = {}
        res["t_recall10"], res["t_precision10"], res["t_measure10"] = \
            mir_eval.hierarchy.tmeasure(ref_hier, est_hier, window=10)
        res["t_recall15"], res["t_precision15"], res["t_measure15"] = \
            mir_eval.hierarchy.tmeasure(ref_hier, est_hier, window=15)

        res["track_id"] = os.path.basename(est_file)[:-5]
        return res
    else:
        # Flat
        return compute_results(ref_inter, est_inter, ref_labels, est_labels,
                               bins, est_file)

def compute_information_gain(ann_inter, est_inter, est_file, bins):
    """Computes the information gain of the est_file from the annotated
    intervals and the estimated intervals."""
    ann_times = utils.intervals_to_times(ann_inter)
    est_times = utils.intervals_to_times(est_inter)
    try:
        D = mir_eval.beat.information_gain(ann_times, est_times, bins=bins)
    except:
        logging.warning("Couldn't compute the Information Gain for file "
                        "%s" % est_file)
        D = 0
    return D

def process_track(file_struct, boundaries_id, labels_id, config,
                  annotator_id=0):
    """Processes a single track.

    Parameters
    ----------
    file_struct : object (FileStruct) or str
        File struct or full path of the audio file to be evaluated.
    boundaries_id : str
        Identifier of the boundaries algorithm.
    labels_id : str
        Identifier of the labels algorithm.
    config : dict
        Configuration of the algorithms to be evaluated.
    annotator_id : int
        Number identifiying the annotator.

    Returns
    -------
    one_res : dict
        Dictionary of the results (see function compute_results).
    """
    # Convert to file_struct if string is passed
    if isinstance(file_struct, six.string_types):
        file_struct = io.FileStruct(file_struct)

    est_file = file_struct.est_file
    ref_file = file_struct.ref_file

    # Sanity check
    assert os.path.basename(est_file)[:-4] == \
        os.path.basename(ref_file)[:-4], "File names are different %s --- %s" \
        % (os.path.basename(est_file)[:-4], os.path.basename(ref_file)[:-4])

    if not os.path.isfile(ref_file):
        raise NoReferencesError("Reference file %s does not exis. You must "
                                "have annotated references to run "
                                "evaluations." % ref_file)

    # TODO: Better exception handling
    try:
        one_res = compute_gt_results(est_file, ref_file, boundaries_id,
                                     labels_id, config,
                                     annotator_id=annotator_id)
    except:
        logging.warning("Could not compute evaluations for %s. Error: %s" %
                        (est_file, sys.exc_info()[1]))
        one_res = []

    return one_res

def get_results_file_name(boundaries_id, labels_id, config, annotator_id):
    """Based on the config and the dataset, get the file name to store the
    results."""
    utils.ensure_dir(msaf.config.results_dir)
    file_name = os.path.join(msaf.config.results_dir, "results")
    file_name += "_boundsE%s_labelsE%s" % (boundaries_id, labels_id)
    file_name += "_annotatorE%d" % (annotator_id)
    sorted_keys = sorted(config.keys(), key=str.lower)
    for key in sorted_keys:
        file_name += "_%sE%s" % (key, str(config[key]).replace("/", "_"))

    # Check for max file length
    if len(file_name) > 255 - len(msaf.config.results_ext):
        file_name = file_name[:255 - len(msaf.config.results_ext)]

    return file_name + msaf.config.results_ext

def process(in_path, boundaries_id=msaf.config.default_bound_id,
            labels_id=msaf.config.default_label_id, annot_beats=False,
            framesync=False, feature="pcp", hier=False, save=False,
            n_jobs=4, annotator_id=0, config=None):
    """Main process.

    Parameters
    ----------
    in_path : str
        Path to the dataset root folder.
    boundaries_id : str
        Boundaries algorithm identifier (e.g. siplca, cnmf)
    labels_id : str
        Labels algorithm identifier (e.g. siplca, cnmf)
    annot_beats : boolean
        Whether to use the annotated beats or not.
    framesync: str
        Whether to use framesync features or not (default: False -> beatsync)
    feature: str
        String representing the feature to be used (e.g. pcp, mfcc, tonnetz)
    hier : bool
        Whether to compute a hierarchical or flat segmentation.
    save: boolean
        Whether to save the results into the SQLite database.
    n_jobs: int
        Number of processes to run in parallel. Only available in collection
        mode.
    annotator_id : int
        Number identifiying the annotator.
    config: dict
        Dictionary containing custom configuration parameters for the
        algorithms.  If None, the default parameters are used.

    Return
    ------
    results : pd.DataFrame
        DataFrame containing the evaluations for each file.
    """

    # Set up configuration based on algorithms parameters
    if config is None:
        config = io.get_configuration(feature, annot_beats, framesync,
                                      boundaries_id, labels_id)

    # Hierarchical segmentation
    config["hier"] = hier

    # Remove actual features
    config.pop("features", None)

    # Get out file in case we want to save results
    out_file = get_results_file_name(boundaries_id, labels_id, config,
                                     annotator_id)

    # All evaluations
    results = pd.DataFrame()

    if os.path.isfile(in_path):
        # Single File mode
        evals = [process_track(in_path, boundaries_id, labels_id, config,
                               annotator_id=annotator_id)]
    else:
        # Collection mode
        # If out_file already exists, do not compute new results
        if os.path.exists(out_file):
            logging.info("Results already exists, reading from file %s" %
                         out_file)
            results = pd.read_csv(out_file)
            print_results(results)
            return results

        # Get files
        file_structs = io.get_dataset_files(in_path)

        logging.info("Evaluating %d tracks..." % len(file_structs))

        # Evaluate in parallel
        evals = Parallel(n_jobs=n_jobs)(delayed(process_track)(
            file_struct, boundaries_id, labels_id, config,
            annotator_id=annotator_id) for file_struct in file_structs[:])

    # Aggregate evaluations in pandas format
    for e in evals:
        if e != []:
            results = results.append(e, ignore_index=True)
    logging.info("%d tracks analyzed" % len(results))

    # Print results
    print_results(results)

    # Save all results
    if save:
        logging.info("Writing results in %s" % out_file)
        results.to_csv(out_file)

    return results
<commit_msg>Only printing averages<commit_after>"""
Evaluates the estimated results of the Segmentation dataset against the
ground truth (human annotated data).
"""
import jams
from joblib import Parallel, delayed
import logging
import mir_eval
import numpy as np
import os
import pandas as pd
import six
import sys

# Local stuff
import msaf
from msaf.exceptions import NoReferencesError
import msaf.input_output as io
from msaf import utils

def print_results(results):
    """Print all the results.

    Parameters
    ----------
    results: pd.DataFrame
        Dataframe with all the results
    """
    if len(results) == 0:
        return
    res = results.mean()
    logging.info("Results:\n%s" % res)

def compute_results(ann_inter, est_inter, ann_labels, est_labels, bins,
                    est_file, weight=0.58):
    """Compute the results using all the available evaluations.

    Parameters
    ----------
    ann_inter : np.array
        Annotated intervals in seconds.
    est_inter : np.array
        Estimated intervals in seconds.
    ann_labels : np.array
        Annotated labels.
    est_labels : np.array
        Estimated labels
    bins : int
        Number of bins for the information gain.
    est_file : str
        Path to the output file to store results.
    weight: float
        Weight the Precision and Recall values of the hit rate boundaries
        differently (<1 will weight Precision higher, >1 will weight Recall
        higher).
        The default parameter (0.58) is the one proposed in (Nieto et al. 2014)

    Return
    ------
    results : dict
        Contains the results of all the evaluations for the given file.
        Keys are the following:
            track_id: Name of the track
            HitRate_3F: F-measure of hit rate at 3 seconds
            HitRate_3P: Precision of hit rate at 3 seconds
            HitRate_3R: Recall of hit rate at 3 seconds
            HitRate_0.5F: F-measure of hit rate at 0.5 seconds
            HitRate_0.5P: Precision of hit rate at 0.5 seconds
            HitRate_0.5R: Recall of hit rate at 0.5 seconds
            HitRate_w3F: F-measure of hit rate at 3 seconds weighted
            HitRate_w0.5F: F-measure of hit rate at 0.5 seconds weighted
            HitRate_wt3F: F-measure of hit rate at 3 seconds weighted and
                          trimmed
            HitRate_wt0.5F: F-measure of hit rate at 0.5 seconds weighted
                            and trimmed
            HitRate_t3F: F-measure of hit rate at 3 seconds (trimmed)
            HitRate_t3P: Precision of hit rate at 3 seconds (trimmed)
            HitRate_t3F: Recall of hit rate at 3 seconds (trimmed)
            HitRate_t0.5F: F-measure of hit rate at 0.5 seconds (trimmed)
            HitRate_t0.5P: Precision of hit rate at 0.5 seconds (trimmed)
            HitRate_t0.5R: Recall of hit rate at 0.5 seconds (trimmed)
            DevA2E: Median deviation of annotation to estimation
            DevE2A: Median deviation of estimation to annotation
            D: Information gain
            PWF: F-measure of pair-wise frame clustering
            PWP: Precision of pair-wise frame clustering
            PWR: Recall of pair-wise frame clustering
            Sf: F-measure normalized entropy score
            So: Oversegmentation normalized entropy score
            Su: Undersegmentation normalized entropy score
    """
    res = {}

    # --Boundaries-- #
    # Hit Rate standard
    res["HitRate_3P"], res["HitRate_3R"], res["HitRate_3F"] = \
        mir_eval.segment.detection(ann_inter, est_inter, window=3, trim=False)
    res["HitRate_0.5P"], res["HitRate_0.5R"], res["HitRate_0.5F"] = \
        mir_eval.segment.detection(ann_inter, est_inter, window=.5, trim=False)

    # Hit rate trimmed
    res["HitRate_t3P"], res["HitRate_t3R"], res["HitRate_t3F"] = \
        mir_eval.segment.detection(ann_inter, est_inter, window=3, trim=True)
    res["HitRate_t0.5P"], res["HitRate_t0.5R"], res["HitRate_t0.5F"] = \
        mir_eval.segment.detection(ann_inter, est_inter, window=.5, trim=True)

    # Hit rate weighted
    res["HitRate_w3P"], _, _ = mir_eval.segment.detection(
        ann_inter, est_inter, window=3, trim=False, beta=weight)
    res["HitRate_w0.5P"], _, _ = mir_eval.segment.detection(
        ann_inter, est_inter, window=.5, trim=False, beta=weight)

    # Hit rate weighted and trimmed
    res["HitRate_wt3P"], _, _ = mir_eval.segment.detection(
        ann_inter, est_inter, window=3, trim=True, beta=weight)
    res["HitRate_wt0.5P"], _, _ = mir_eval.segment.detection(
        ann_inter, est_inter, window=.5, trim=True, beta=weight)

    # Information gain
    res["D"] = compute_information_gain(ann_inter, est_inter, est_file,
                                        bins=bins)

    # Median Deviations
    res["DevR2E"], res["DevE2R"] = mir_eval.segment.deviation(
        ann_inter, est_inter, trim=False)
    res["DevtR2E"], res["DevtE2R"] = mir_eval.segment.deviation(
        ann_inter, est_inter, trim=True)

    # --Labels-- #
    if est_labels is not None and "-1" in est_labels:
        est_labels = None

    if est_labels is not None and len(est_labels) != 0:
        try:
            # Align labels with intervals
            ann_labels = list(ann_labels)
            est_labels = list(est_labels)
            ann_inter, ann_labels = mir_eval.util.adjust_intervals(ann_inter,
                                                                   ann_labels)
            est_inter, est_labels = mir_eval.util.adjust_intervals(
                est_inter, est_labels, t_min=0, t_max=ann_inter.max())

            # Pair-wise frame clustering
            res["PWP"], res["PWR"], res["PWF"] = mir_eval.segment.pairwise(
                ann_inter, ann_labels, est_inter, est_labels)

            # Normalized Conditional Entropies
            res["So"], res["Su"], res["Sf"] = mir_eval.segment.nce(
                ann_inter, ann_labels, est_inter, est_labels)
        except:
            logging.warning("Labeling evaluation failed in file: %s" %
                            est_file)
            return {}

    # Names
    base = os.path.basename(est_file)
    res["track_id"] = base[:-5]

    return res

def compute_gt_results(est_file, ref_file, boundaries_id, labels_id, config,
                       bins=251, annotator_id=0):
    """Computes the results by using the ground truth dataset identified by
    the annotator parameter.

    Return
    ------
    results : dict
        Dictionary of the results (see function compute_results).
    """
    try:
        if config["hier"]:
            ref_times, ref_labels, ref_levels = \
                msaf.io.read_hier_references(
                    ref_file, annotation_id=0,
                    exclude_levels=["segment_salami_function"])
        else:
            jam = jams.load(ref_file, validate=False)
            ann = jam.search(namespace='segment_.*')[annotator_id]
            ref_inter, ref_labels = ann.data.to_interval_values()
    except:
        logging.warning("No references for file: %s" % ref_file)
        return {}

    # Read estimations with correct configuration
    est_inter, est_labels = io.read_estimations(est_file, boundaries_id,
                                                labels_id, **config)
    if len(est_inter) == 0:
        logging.warning("No estimations for file: %s" % est_file)
        return {}

    # Compute the results and return
    logging.info("Evaluating %s" % os.path.basename(est_file))
    if config["hier"]:
        # Hierarchical
        assert len(est_inter) == len(est_labels), "Same number of levels " \
            "are required in the boundaries and labels for the hierarchical " \
            "evaluation."
        est_times = []
        est_labels = []

        # Sort based on how many segments per level
        est_inter = sorted(est_inter, key=lambda level: len(level))

        for inter in est_inter:
            est_times.append(msaf.utils.intervals_to_times(inter))
            # Add fake labels (hierarchical eval does not use labels --yet--)
            est_labels.append(np.ones(len(est_times[-1]) - 1) * -1)

        # Align the times
        utils.align_end_hierarchies(est_times, ref_times, thres=1)

        # To intervals
        est_hier = [utils.times_to_intervals(times) for times in est_times]
        ref_hier = [utils.times_to_intervals(times) for times in ref_times]

        # Compute evaluations
        res = {}
        res["t_recall10"], res["t_precision10"], res["t_measure10"] = \
            mir_eval.hierarchy.tmeasure(ref_hier, est_hier, window=10)
        res["t_recall15"], res["t_precision15"], res["t_measure15"] = \
            mir_eval.hierarchy.tmeasure(ref_hier, est_hier, window=15)

        res["track_id"] = os.path.basename(est_file)[:-5]
        return res
    else:
        # Flat
        return compute_results(ref_inter, est_inter, ref_labels, est_labels,
                               bins, est_file)

def compute_information_gain(ann_inter, est_inter, est_file, bins):
    """Computes the information gain of the est_file from the annotated
    intervals and the estimated intervals."""
    ann_times = utils.intervals_to_times(ann_inter)
    est_times = utils.intervals_to_times(est_inter)
    try:
        D = mir_eval.beat.information_gain(ann_times, est_times, bins=bins)
    except:
        logging.warning("Couldn't compute the Information Gain for file "
                        "%s" % est_file)
        D = 0
    return D

def process_track(file_struct, boundaries_id, labels_id, config,
                  annotator_id=0):
    """Processes a single track.

    Parameters
    ----------
    file_struct : object (FileStruct) or str
        File struct or full path of the audio file to be evaluated.
    boundaries_id : str
        Identifier of the boundaries algorithm.
    labels_id : str
        Identifier of the labels algorithm.
    config : dict
        Configuration of the algorithms to be evaluated.
    annotator_id : int
        Number identifiying the annotator.

    Returns
    -------
    one_res : dict
        Dictionary of the results (see function compute_results).
    """
    # Convert to file_struct if string is passed
    if isinstance(file_struct, six.string_types):
        file_struct = io.FileStruct(file_struct)

    est_file = file_struct.est_file
    ref_file = file_struct.ref_file

    # Sanity check
    assert os.path.basename(est_file)[:-4] == \
        os.path.basename(ref_file)[:-4], "File names are different %s --- %s" \
        % (os.path.basename(est_file)[:-4], os.path.basename(ref_file)[:-4])

    if not os.path.isfile(ref_file):
        raise NoReferencesError("Reference file %s does not exis. You must "
                                "have annotated references to run "
                                "evaluations." % ref_file)

    # TODO: Better exception handling
    try:
        one_res = compute_gt_results(est_file, ref_file, boundaries_id,
                                     labels_id, config,
                                     annotator_id=annotator_id)
    except:
        logging.warning("Could not compute evaluations for %s. Error: %s" %
                        (est_file, sys.exc_info()[1]))
        one_res = []

    return one_res

def get_results_file_name(boundaries_id, labels_id, config, annotator_id):
    """Based on the config and the dataset, get the file name to store the
    results."""
    utils.ensure_dir(msaf.config.results_dir)
    file_name = os.path.join(msaf.config.results_dir, "results")
    file_name += "_boundsE%s_labelsE%s" % (boundaries_id, labels_id)
    file_name += "_annotatorE%d" % (annotator_id)
    sorted_keys = sorted(config.keys(), key=str.lower)
    for key in sorted_keys:
        file_name += "_%sE%s" % (key, str(config[key]).replace("/", "_"))

    # Check for max file length
    if len(file_name) > 255 - len(msaf.config.results_ext):
        file_name = file_name[:255 - len(msaf.config.results_ext)]

    return file_name + msaf.config.results_ext

def process(in_path, boundaries_id=msaf.config.default_bound_id,
            labels_id=msaf.config.default_label_id, annot_beats=False,
            framesync=False, feature="pcp", hier=False, save=False,
            n_jobs=4, annotator_id=0, config=None):
    """Main process.

    Parameters
    ----------
    in_path : str
        Path to the dataset root folder.
    boundaries_id : str
        Boundaries algorithm identifier (e.g. siplca, cnmf)
    labels_id : str
        Labels algorithm identifier (e.g. siplca, cnmf)
    annot_beats : boolean
        Whether to use the annotated beats or not.
    framesync: str
        Whether to use framesync features or not (default: False -> beatsync)
    feature: str
        String representing the feature to be used (e.g. pcp, mfcc, tonnetz)
    hier : bool
        Whether to compute a hierarchical or flat segmentation.
    save: boolean
        Whether to save the results into the SQLite database.
    n_jobs: int
        Number of processes to run in parallel. Only available in collection
        mode.
    annotator_id : int
        Number identifiying the annotator.
    config: dict
        Dictionary containing custom configuration parameters for the
        algorithms.  If None, the default parameters are used.

    Return
    ------
    results : pd.DataFrame
        DataFrame containing the evaluations for each file.
    """

    # Set up configuration based on algorithms parameters
    if config is None:
        config = io.get_configuration(feature, annot_beats, framesync,
                                      boundaries_id, labels_id)

    # Hierarchical segmentation
    config["hier"] = hier

    # Remove actual features
    config.pop("features", None)

    # Get out file in case we want to save results
    out_file = get_results_file_name(boundaries_id, labels_id, config,
                                     annotator_id)

    # All evaluations
    results = pd.DataFrame()

    if os.path.isfile(in_path):
        # Single File mode
        evals = [process_track(in_path, boundaries_id, labels_id, config,
                               annotator_id=annotator_id)]
    else:
        # Collection mode
        # If out_file already exists, do not compute new results
        if os.path.exists(out_file):
            logging.info("Results already exists, reading from file %s" %
                         out_file)
            results = pd.read_csv(out_file)
            print_results(results)
            return results

        # Get files
        file_structs = io.get_dataset_files(in_path)

        logging.info("Evaluating %d tracks..." % len(file_structs))

        # Evaluate in parallel
        evals = Parallel(n_jobs=n_jobs)(delayed(process_track)(
            file_struct, boundaries_id, labels_id, config,
            annotator_id=annotator_id) for file_struct in file_structs[:])

    # Aggregate evaluations in pandas format
    for e in evals:
        if e != []:
            results = results.append(e, ignore_index=True)
    logging.info("%d tracks analyzed" % len(results))

    # Print results
    print_results(results)

    # Save all results
    if save:
        logging.info("Writing results in %s" % out_file)
        results.to_csv(out_file)

    return results
<|endoftext|>