<commit_before>from __future__ import print_function, division

import numpy as np

import yt.frontends.stream.api as stream
from yt.utilities.io_handler import BaseIOHandler

class HyperionIOHandler(BaseIOHandler):

    def __init__(self, grids, dust_id):
        if np.array(grids[0][grids[0].keys()[0]]).ndim == 4:
            self.dust_id = dust_id
        else:
            self.dust_id = None
        self.grids = grids
        BaseIOHandler.__init__(self)

    def _read_data_set(self, grid, field):
        if self.dust_id is None:
            return np.array(self.grids[grid.id][field.lower()].transpose())
        else:
            return np.array(self.grids[grid.id][field.lower()][self.dust_id].transpose())

    def _read_data_slice(self, grid, field, axis, coord):
        sl = [slice(None), slice(None), slice(None)]
        sl[axis] = slice(coord, coord + 1)
        return self._read_data_set(grid, field)[sl]

class StreamFieldData(object):

    def __init__(self, fields):
        self.fields = fields

    @property
    def all_fields(self):
        return self.fields

def almost_equal(a, b):
    return a / b < 1. + 1.e-4 and b / a < 1. + 1.e-4

def amr_grid_to_yt_stream(levels, dust_id=0):

    # Try and guess the refinement ratio - if it is not constant, then
    # we can't use yt

    if len(levels) == 0 or len(levels[0].grids) == 0:
        raise Exception("Need at least one level with one grid to convert to a yt object")
    elif len(levels) == 1:
        refine = 2
    else:
        dx = []
        dy = []
        dz = []
        for ilevel, level in enumerate(levels):
            for igrid, grid in enumerate(level.grids):
                gdx = (grid.xmax - grid.xmin) / float(grid.nx)
                gdy = (grid.ymax - grid.ymin) / float(grid.ny)
                gdz = (grid.zmax - grid.zmin) / float(grid.nz)
                if igrid == 0:
                    dx.append(gdx)
                    dy.append(gdy)
                    dz.append(gdz)
                else:
                    if not almost_equal(dx[-1], gdx):
                        raise Exception("dx scale differs between grids in level %i (expected %g and got %g)" % (ilevel, dx[-1], gdx))
                    if not almost_equal(dy[-1], gdy):
                        raise Exception("dy scale differs between grids in level %i (expected %g and got %g)" % (ilevel, dy[-1], gdy))
                    if not almost_equal(dz[-1], gdz):
                        raise Exception("dz scale differs between grids in level %i (expected %g and got %g)" % (ilevel, dz[-1], gdz))
        dx = np.array(dx)
        dy = np.array(dy)
        dz = np.array(dz)
        refine_x = dx[:-1] / dx[1:]
        refine_y = dy[:-1] / dy[1:]
        refine_z = dz[:-1] / dz[1:]
        for i in range(len(levels) - 1):
            if abs(refine_x[i] - round(refine_x[i])) > 1.e-5:
                raise Exception("refinement ratio is not an integer (%g)" % refine_x[i])
            if abs(refine_y[i] - round(refine_y[i])) > 1.e-5:
                raise Exception("refinement ratio is not an integer (%g)" % refine_y[i])
            if abs(refine_z[i] - round(refine_z[i])) > 1.e-5:
                raise Exception("refinement ratio is not an integer (%g)" % refine_z[i])
        refine_x = np.round(refine_x).astype(int)
        refine_y = np.round(refine_y).astype(int)
        refine_z = np.round(refine_z).astype(int)
        if not np.all(np.hstack([refine_x, refine_y, refine_z]) == refine_x[0]):
            raise Exception("refinement ratio changes between levels and/or directions (x = %s, y = %s, z = %s)" % (str(refine_x), str(refine_y), str(refine_z)))
        refine = int(refine_x[0])

    xmin = ymin = zmin = +np.inf
    xmax = ymax = zmax = -np.inf

    left_edge = []
    right_edge = []
    dimensions = []
    level_ids = []

    n_grids = 0
    grids = []

    for ilevel, level in enumerate(levels):

        for grid in level.grids:

            left_edge.append([grid.xmin,
                              grid.ymin,
                              grid.zmin])

            right_edge.append([grid.xmax,
                               grid.ymax,
                               grid.zmax])

            dimensions.append([grid.nx,
                               grid.ny,
                               grid.nz])

            level_ids.append(ilevel)

            n_grids += 1

            xmin = min(xmin, grid.xmin)
            xmax = max(xmax, grid.xmax)
            ymin = min(ymin, grid.ymin)
            ymax = max(ymax, grid.ymax)
            zmin = min(zmin, grid.zmin)
            zmax = max(zmax, grid.zmax)

            grids.append(grid.quantities)

    left_edge = np.array(left_edge, dtype=np.float64)
    right_edge = np.array(right_edge, dtype=np.float64)
    dimensions = np.array(dimensions, dtype=np.int32)

    # The reshape is necessary due to a strange Numpy bug
    level_ids = np.array(level_ids, dtype=np.int32).reshape((len(level_ids), 1))

    parent_ids = None
    particle_count = np.zeros((n_grids, 1), dtype='int32')
    processor_ids = np.zeros(n_grids)

    grid = levels[0].grids[0]

    # Determine domain resolution

    dx = (grid.xmax - grid.xmin) / float(grid.nx)
    nx = int(round((xmax - xmin) / dx))

    dy = (grid.ymax - grid.ymin) / float(grid.ny)
    ny = int(round((ymax - ymin) / dy))

    dz = (grid.zmax - grid.zmin) / float(grid.nz)
    nz = int(round((zmax - zmin) / dz))

    # Determine fields

    fields = grid.quantities.keys()

    # Set up StreamHandler

    handler = stream.StreamHandler(
        left_edge[:],
        right_edge[:],
        dimensions[:],
        level_ids[:],
        parent_ids,
        particle_count[:],
        processor_ids[:],
        StreamFieldData(fields),
        HyperionIOHandler(grids, dust_id),
    )

    handler.name = 'hyperion'
    handler.domain_left_edge = np.array([xmin, ymin, zmin])
    handler.domain_right_edge = np.array([xmax, ymax, zmax])
    handler.refine_by = refine
    handler.dimensionality = 3
    handler.domain_dimensions = np.array([nx, ny, nz])
    handler.simulation_time = 0.0
    handler.cosmology_simulation = 0

    spf = stream.StreamStaticOutput(handler)
    spf.units["cm"] = 1.0
    spf.units["unitary"] = 1.0 / ((spf.domain_right_edge - spf.domain_left_edge).max())

    return spf

def edge_list(refined, xmin, xmax, ymin, ymax, zmin, zmax, i=0):

    if not refined[i]:
        return [], i
    else:
        xmid = (xmin + xmax) * 0.5
        ymid = (ymin + ymax) * 0.5
        zmid = (zmin + zmax) * 0.5
        e1, i = edge_list(refined, xmin, xmid, ymin, ymid, zmin, zmid, i=i + 1)
        e2, i = edge_list(refined, xmid, xmax, ymin, ymid, zmin, zmid, i=i + 1)
        e3, i = edge_list(refined, xmin, xmid, ymid, ymax, zmin, zmid, i=i + 1)
        e4, i = edge_list(refined, xmid, xmax, ymid, ymax, zmin, zmid, i=i + 1)
        e5, i = edge_list(refined, xmin, xmid, ymin, ymid, zmid, zmax, i=i + 1)
        e6, i = edge_list(refined, xmid, xmax, ymin, ymid, zmid, zmax, i=i + 1)
        e7, i = edge_list(refined, xmin, xmid, ymid, ymax, zmid, zmax, i=i + 1)
        e8, i = edge_list(refined, xmid, xmax, ymid, ymax, zmid, zmax, i=i + 1)
        return [(xmin, xmax, ymin, ymax, zmin, zmax)] + e1 + e2 + e3 + e4 + e5 + e6 + e7 + e8, i

def decompose_quantity(refined, array, i=0):
    subarray = []
    newarray = []
    if refined[i]:
        for s in range(8):
            i = i + 1
            subarray.append(array[i])
            a, i = decompose_quantity(refined, array, i=i)
            newarray += a
        subarray = np.array(subarray).reshape((2, 2, 2))
        return [subarray] + newarray, i
    else:
        return [], i

def level_list(refined, i=0, level=0):
    if refined[i]:
        levels = [level]
        for s in range(8):
            l, i = level_list(refined, i=i + 1, level=level + 1)
            levels = levels + l
        return levels, i
    else:
        return [], i

class HyperionIOHandlerOct(BaseIOHandler):

    def __init__(self, quantities):
        self.quantities = quantities
        BaseIOHandler.__init__(self)

    def _read_data_set(self, grid, field):
        return self.quantities[field.lower()][grid.id].transpose()

    def _read_data_slice(self, grid, field, axis, coord):
        sl = [slice(None), slice(None), slice(None)]
        sl[axis] = slice(coord, coord + 1)
        return self._read_data_set(grid, field)[sl]

def octree_grid_to_yt_stream(grid, dust_id=0):

    xmin = grid.x - grid.dx
    xmax = grid.x + grid.dx
    ymin = grid.y - grid.dy
    ymax = grid.y + grid.dy
    zmin = grid.z - grid.dz
    zmax = grid.z + grid.dz

    e, i = edge_list(grid.refined, xmin, xmax, ymin, ymax, zmin, zmax)

    e = np.array(e)

    left_edge = e[:, ::2].astype(np.float64)
    right_edge = e[:, 1::2].astype(np.float64)
    dimensions = np.ones((np.sum(grid.refined), 3), dtype=np.int32) * 2
    level_ids = np.array(level_list(grid.refined)[0], dtype=np.int32)
    level_ids = level_ids.reshape((len(level_ids), 1))

    n_grids = np.sum(grid.refined)

    parent_ids = None
    particle_count = np.zeros((n_grids, 1), dtype='int32')
    processor_ids = np.zeros(n_grids)

    # Determine fields

    fields = grid.quantities.keys()

    # Set up StreamHandler

    quantities = {}
    for field in grid.quantities:
        quantities[field] = decompose_quantity(grid.refined, grid.quantities[field][dust_id])[0]

    handler = stream.StreamHandler(
        left_edge[:],
        right_edge[:],
        dimensions[:],
        level_ids[:],
        parent_ids,
        particle_count[:],
        processor_ids[:],
        StreamFieldData(fields),
        HyperionIOHandlerOct(quantities),
    )

    handler.name = 'hyperion'
    handler.domain_left_edge = np.array([xmin, ymin, zmin])
    handler.domain_right_edge = np.array([xmax, ymax, zmax])
    handler.refine_by = 2
    handler.dimensionality = 3
    handler.domain_dimensions = np.array([2, 2, 2])
    handler.simulation_time = 0.0
    handler.cosmology_simulation = 0

    spf = stream.StreamStaticOutput(handler)
    spf.units["cm"] = 1.0
    spf.units["unitary"] = 1.0 / ((spf.domain_right_edge - spf.domain_left_edge).max())

    return spf

def cartesian_grid_to_yt_stream(grid, xmin, xmax, ymin, ymax, zmin, zmax, dust_id=0):

    nz, ny, nx = grid.shape

    grids = [grid.quantities]
    left_edge = np.array([[xmin, ymin, zmin]], dtype=np.float64)
    right_edge = np.array([[xmax, ymax, zmax]], dtype=np.float64)
    dimensions = np.array([[nx, ny, nz]], dtype=np.int32)
    level_ids = np.array([[0]], dtype=np.int32).reshape((1, 1))

    parent_ids = None
    particle_count = np.zeros((1, 1), dtype='int32')
    processor_ids = np.zeros(1)

    # Determine fields

    fields = grid.quantities.keys()

    # Set up StreamHandler

    handler = stream.StreamHandler(
        left_edge[:],
        right_edge[:],
        dimensions[:],
        level_ids[:],
        parent_ids,
        particle_count[:],
        processor_ids[:],
        StreamFieldData(fields),
        HyperionIOHandler(grids, dust_id),
    )

    handler.name = 'hyperion'
    handler.domain_left_edge = np.array([xmin, ymin, zmin])
    handler.domain_right_edge = np.array([xmax, ymax, zmax])
    handler.refine_by = 2
    handler.dimensionality = 3
    handler.domain_dimensions = np.array([nx, ny, nz])
    handler.simulation_time = 0.0
    handler.cosmology_simulation = 0

    spf = stream.StreamStaticOutput(handler)
    spf.units["cm"] = 1.0
    spf.units["unitary"] = 1.0 / ((spf.domain_right_edge - spf.domain_left_edge).max())

    return spf<commit_msg>[yt wrappers] Ensuring units logic matches logic in yt's stream frontend.<commit_after>from __future__ import print_function, division

import numpy as np

import yt.frontends.stream.api as stream
from yt.utilities.io_handler import BaseIOHandler
from yt.utilities.definitions import mpc_conversion

class HyperionIOHandler(BaseIOHandler):

    def __init__(self, grids, dust_id):
        if np.array(grids[0][grids[0].keys()[0]]).ndim == 4:
            self.dust_id = dust_id
        else:
            self.dust_id = None
        self.grids = grids
        BaseIOHandler.__init__(self)

    def _read_data_set(self, grid, field):
        if self.dust_id is None:
            return np.array(self.grids[grid.id][field.lower()].transpose())
        else:
            return np.array(self.grids[grid.id][field.lower()][self.dust_id].transpose())

    def _read_data_slice(self, grid, field, axis, coord):
        sl = [slice(None), slice(None), slice(None)]
        sl[axis] = slice(coord, coord + 1)
        return self._read_data_set(grid, field)[sl]

class StreamFieldData(object):

    def __init__(self, fields):
        self.fields = fields

    @property
    def all_fields(self):
        return self.fields

def almost_equal(a, b):
    return a / b < 1. + 1.e-4 and b / a < 1. + 1.e-4

def amr_grid_to_yt_stream(levels, dust_id=0):

    # Try and guess the refinement ratio - if it is not constant, then
    # we can't use yt

    if len(levels) == 0 or len(levels[0].grids) == 0:
        raise Exception("Need at least one level with one grid to convert to a yt object")
    elif len(levels) == 1:
        refine = 2
    else:
        dx = []
        dy = []
        dz = []
        for ilevel, level in enumerate(levels):
            for igrid, grid in enumerate(level.grids):
                gdx = (grid.xmax - grid.xmin) / float(grid.nx)
                gdy = (grid.ymax - grid.ymin) / float(grid.ny)
                gdz = (grid.zmax - grid.zmin) / float(grid.nz)
                if igrid == 0:
                    dx.append(gdx)
                    dy.append(gdy)
                    dz.append(gdz)
                else:
                    if not almost_equal(dx[-1], gdx):
                        raise Exception("dx scale differs between grids in level %i (expected %g and got %g)" % (ilevel, dx[-1], gdx))
                    if not almost_equal(dy[-1], gdy):
                        raise Exception("dy scale differs between grids in level %i (expected %g and got %g)" % (ilevel, dy[-1], gdy))
                    if not almost_equal(dz[-1], gdz):
                        raise Exception("dz scale differs between grids in level %i (expected %g and got %g)" % (ilevel, dz[-1], gdz))
        dx = np.array(dx)
        dy = np.array(dy)
        dz = np.array(dz)
        refine_x = dx[:-1] / dx[1:]
        refine_y = dy[:-1] / dy[1:]
        refine_z = dz[:-1] / dz[1:]
        for i in range(len(levels) - 1):
            if abs(refine_x[i] - round(refine_x[i])) > 1.e-5:
                raise Exception("refinement ratio is not an integer (%g)" % refine_x[i])
            if abs(refine_y[i] - round(refine_y[i])) > 1.e-5:
                raise Exception("refinement ratio is not an integer (%g)" % refine_y[i])
            if abs(refine_z[i] - round(refine_z[i])) > 1.e-5:
                raise Exception("refinement ratio is not an integer (%g)" % refine_z[i])
        refine_x = np.round(refine_x).astype(int)
        refine_y = np.round(refine_y).astype(int)
        refine_z = np.round(refine_z).astype(int)
        if not np.all(np.hstack([refine_x, refine_y, refine_z]) == refine_x[0]):
            raise Exception("refinement ratio changes between levels and/or directions (x = %s, y = %s, z = %s)" % (str(refine_x), str(refine_y), str(refine_z)))
        refine = int(refine_x[0])

    xmin = ymin = zmin = +np.inf
    xmax = ymax = zmax = -np.inf

    left_edge = []
    right_edge = []
    dimensions = []
    level_ids = []

    n_grids = 0
    grids = []

    for ilevel, level in enumerate(levels):

        for grid in level.grids:

            left_edge.append([grid.xmin,
                              grid.ymin,
                              grid.zmin])

            right_edge.append([grid.xmax,
                               grid.ymax,
                               grid.zmax])

            dimensions.append([grid.nx,
                               grid.ny,
                               grid.nz])

            level_ids.append(ilevel)

            n_grids += 1

            xmin = min(xmin, grid.xmin)
            xmax = max(xmax, grid.xmax)
            ymin = min(ymin, grid.ymin)
            ymax = max(ymax, grid.ymax)
            zmin = min(zmin, grid.zmin)
            zmax = max(zmax, grid.zmax)

            grids.append(grid.quantities)

    left_edge = np.array(left_edge, dtype=np.float64)
    right_edge = np.array(right_edge, dtype=np.float64)
    dimensions = np.array(dimensions, dtype=np.int32)

    # The reshape is necessary due to a strange Numpy bug
    level_ids = np.array(level_ids, dtype=np.int32).reshape((len(level_ids), 1))

    parent_ids = None
    particle_count = np.zeros((n_grids, 1), dtype='int32')
    processor_ids = np.zeros(n_grids)

    grid = levels[0].grids[0]

    # Determine domain resolution

    dx = (grid.xmax - grid.xmin) / float(grid.nx)
    nx = int(round((xmax - xmin) / dx))

    dy = (grid.ymax - grid.ymin) / float(grid.ny)
    ny = int(round((ymax - ymin) / dy))

    dz = (grid.zmax - grid.zmin) / float(grid.nz)
    nz = int(round((zmax - zmin) / dz))

    # Determine fields

    fields = grid.quantities.keys()

    # Set up StreamHandler

    handler = stream.StreamHandler(
        left_edge[:],
        right_edge[:],
        dimensions[:],
        level_ids[:],
        parent_ids,
        particle_count[:],
        processor_ids[:],
        StreamFieldData(fields),
        HyperionIOHandler(grids, dust_id),
    )

    handler.name = 'hyperion'
    handler.domain_left_edge = np.array([xmin, ymin, zmin])
    handler.domain_right_edge = np.array([xmax, ymax, zmax])
    handler.refine_by = refine
    handler.dimensionality = 3
    handler.domain_dimensions = np.array([nx, ny, nz])
    handler.simulation_time = 0.0
    handler.cosmology_simulation = 0

    spf = stream.StreamStaticOutput(handler)
    spf.units["cm"] = 1.0
    spf.units["unitary"] = 1.0 / ((spf.domain_right_edge - spf.domain_left_edge).max())
    spf.units['1'] = 1.0
    box_in_mpc = 1.0 / mpc_conversion['cm']
    for unit in mpc_conversion.keys():
        spf.units[unit] = mpc_conversion[unit] * box_in_mpc

    return spf

def edge_list(refined, xmin, xmax, ymin, ymax, zmin, zmax, i=0):

    if not refined[i]:
        return [], i
    else:
        xmid = (xmin + xmax) * 0.5
        ymid = (ymin + ymax) * 0.5
        zmid = (zmin + zmax) * 0.5
        e1, i = edge_list(refined, xmin, xmid, ymin, ymid, zmin, zmid, i=i + 1)
        e2, i = edge_list(refined, xmid, xmax, ymin, ymid, zmin, zmid, i=i + 1)
        e3, i = edge_list(refined, xmin, xmid, ymid, ymax, zmin, zmid, i=i + 1)
        e4, i = edge_list(refined, xmid, xmax, ymid, ymax, zmin, zmid, i=i + 1)
        e5, i = edge_list(refined, xmin, xmid, ymin, ymid, zmid, zmax, i=i + 1)
        e6, i = edge_list(refined, xmid, xmax, ymin, ymid, zmid, zmax, i=i + 1)
        e7, i = edge_list(refined, xmin, xmid, ymid, ymax, zmid, zmax, i=i + 1)
        e8, i = edge_list(refined, xmid, xmax, ymid, ymax, zmid, zmax, i=i + 1)
        return [(xmin, xmax, ymin, ymax, zmin, zmax)] + e1 + e2 + e3 + e4 + e5 + e6 + e7 + e8, i

def decompose_quantity(refined, array, i=0):
    subarray = []
    newarray = []
    if refined[i]:
        for s in range(8):
            i = i + 1
            subarray.append(array[i])
            a, i = decompose_quantity(refined, array, i=i)
            newarray += a
        subarray = np.array(subarray).reshape((2, 2, 2))
        return [subarray] + newarray, i
    else:
        return [], i

def level_list(refined, i=0, level=0):
    if refined[i]:
        levels = [level]
        for s in range(8):
            l, i = level_list(refined, i=i + 1, level=level + 1)
            levels = levels + l
        return levels, i
    else:
        return [], i

class HyperionIOHandlerOct(BaseIOHandler):

    def __init__(self, quantities):
        self.quantities = quantities
        BaseIOHandler.__init__(self)

    def _read_data_set(self, grid, field):
        return self.quantities[field.lower()][grid.id].transpose()

    def _read_data_slice(self, grid, field, axis, coord):
        sl = [slice(None), slice(None), slice(None)]
        sl[axis] = slice(coord, coord + 1)
        return self._read_data_set(grid, field)[sl]

def octree_grid_to_yt_stream(grid, dust_id=0):

    xmin = grid.x - grid.dx
    xmax = grid.x + grid.dx
    ymin = grid.y - grid.dy
    ymax = grid.y + grid.dy
    zmin = grid.z - grid.dz
    zmax = grid.z + grid.dz

    e, i = edge_list(grid.refined, xmin, xmax, ymin, ymax, zmin, zmax)

    e = np.array(e)

    left_edge = e[:, ::2].astype(np.float64)
    right_edge = e[:, 1::2].astype(np.float64)
    dimensions = np.ones((np.sum(grid.refined), 3), dtype=np.int32) * 2
    level_ids = np.array(level_list(grid.refined)[0], dtype=np.int32)
    level_ids = level_ids.reshape((len(level_ids), 1))

    n_grids = np.sum(grid.refined)

    parent_ids = None
    particle_count = np.zeros((n_grids, 1), dtype='int32')
    processor_ids = np.zeros(n_grids)

    # Determine fields

    fields = grid.quantities.keys()

    # Set up StreamHandler

    quantities = {}
    for field in grid.quantities:
        quantities[field] = decompose_quantity(grid.refined, grid.quantities[field][dust_id])[0]

    handler = stream.StreamHandler(
        left_edge[:],
        right_edge[:],
        dimensions[:],
        level_ids[:],
        parent_ids,
        particle_count[:],
        processor_ids[:],
        StreamFieldData(fields),
        HyperionIOHandlerOct(quantities),
    )

    handler.name = 'hyperion'
    handler.domain_left_edge = np.array([xmin, ymin, zmin])
    handler.domain_right_edge = np.array([xmax, ymax, zmax])
    handler.refine_by = 2
    handler.dimensionality = 3
    handler.domain_dimensions = np.array([2, 2, 2])
    handler.simulation_time = 0.0
    handler.cosmology_simulation = 0

    spf = stream.StreamStaticOutput(handler)
    spf.units["cm"] = 1.0
    spf.units["unitary"] = 1.0 / ((spf.domain_right_edge - spf.domain_left_edge).max())
    spf.units['1'] = 1.0
    box_in_mpc = 1.0 / mpc_conversion['cm']
    for unit in mpc_conversion.keys():
        spf.units[unit] = mpc_conversion[unit] * box_in_mpc

    return spf

def cartesian_grid_to_yt_stream(grid, xmin, xmax, ymin, ymax, zmin, zmax, dust_id=0):

    nz, ny, nx = grid.shape

    grids = [grid.quantities]
    left_edge = np.array([[xmin, ymin, zmin]], dtype=np.float64)
    right_edge = np.array([[xmax, ymax, zmax]], dtype=np.float64)
    dimensions = np.array([[nx, ny, nz]], dtype=np.int32)
    level_ids = np.array([[0]], dtype=np.int32).reshape((1, 1))

    parent_ids = None
    particle_count = np.zeros((1, 1), dtype='int32')
    processor_ids = np.zeros(1)

    # Determine fields

    fields = grid.quantities.keys()

    # Set up StreamHandler

    handler = stream.StreamHandler(
        left_edge[:],
        right_edge[:],
        dimensions[:],
        level_ids[:],
        parent_ids,
        particle_count[:],
        processor_ids[:],
        StreamFieldData(fields),
        HyperionIOHandler(grids, dust_id),
    )

    handler.name = 'hyperion'
    handler.domain_left_edge = np.array([xmin, ymin, zmin])
    handler.domain_right_edge = np.array([xmax, ymax, zmax])
    handler.refine_by = 2
    handler.dimensionality = 3
    handler.domain_dimensions = np.array([nx, ny, nz])
    handler.simulation_time = 0.0
    handler.cosmology_simulation = 0

    spf = stream.StreamStaticOutput(handler)
    spf.units["cm"] = 1.0
    spf.units["unitary"] = 1.0 / ((spf.domain_right_edge - spf.domain_left_edge).max())
    spf.units['1'] = 1.0
    box_in_mpc = 1.0 / mpc_conversion['cm']
    for unit in mpc_conversion.keys():
        spf.units[unit] = mpc_conversion[unit] * box_in_mpc

    return spf<|endoftext|>


<commit_before># Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from pants.backend.go.lint.vet import skip_field
from pants.backend.go.lint.vet.rules import rules as go_vet_rules

def rules():
    return (
        *go_vet_rules(),
        *skip_field.rules(),
    )
<commit_msg>[internal] fix non-empty __init__.py (#13730)<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import os
import io
import json
import pytest

from boussole.conf.model import Settings

def test_001(compiler, temp_builds_dir):
    """compiler.SassCompileHelper.safe_compile: Basic sample without source map"""
    basedir = temp_builds_dir.join('compiler_safecompile_001').strpath

    basic_settings = Settings(initial={
        'SOURCES_PATH': 'scss',
        'TARGET_PATH': 'css',
        'SOURCE_MAP': False,
        'OUTPUT_STYLES': 'compact',
    })

    sourcedir = os.path.join(basedir, basic_settings.SOURCES_PATH)
    targetdir = os.path.join(basedir, basic_settings.TARGET_PATH)

    src = os.path.join(sourcedir, "app.scss")
    dst = os.path.join(targetdir, "app.css")
    src_map = os.path.join(targetdir, "app.map")

    os.makedirs(sourcedir)
    os.makedirs(targetdir)

    # Create sample source to compile
    with io.open(src, 'w', encoding='utf-8') as f:
        result = f.write(
            u"""#content{
            color: #ff0000;
            font-weight: bold;
            &.foo{ border: 1px solid #000000; }
            }
            """)

    # Compile
    success, message = compiler.safe_compile(basic_settings, src, dst)

    assert os.path.exists(dst) == True
    assert os.path.exists(src_map) == False

    # Assert compiled file is ok
    with io.open(dst, 'r', encoding='utf-8') as f:
        content = f.read()

    assert content == ("""#content { color: #ff0000; font-weight: bold; }\n\n"""
                      """#content.foo { border: 1px solid #000000; }\n""")

def test_002(compiler, temp_builds_dir):
    """compiler.SassCompileHelper.safe_compile: Same as basic sample but with
       source map"""
    basedir = temp_builds_dir.join('compiler_safecompile_002').strpath

    basic_settings = Settings(initial={
        'SOURCES_PATH': '.',
        'TARGET_PATH': 'css',
        'SOURCE_MAP': True,
        'OUTPUT_STYLES': 'compact',
    })

    sourcedir = os.path.normpath(os.path.join(basedir, basic_settings.SOURCES_PATH))
    targetdir = os.path.normpath(os.path.join(basedir, basic_settings.TARGET_PATH))

    os.makedirs(sourcedir)
    os.makedirs(targetdir)

    src = os.path.join(sourcedir, "app.scss")
    dst = os.path.join(targetdir, "app.css")
    src_map = os.path.join(targetdir, "app.map")

    # Create sample source to compile
    with io.open(src, 'w', encoding='utf-8') as f:
        result = f.write(u"""#content{ color:#ff0000; font-weight:bold; }""")

    # Compile
    success, message = compiler.safe_compile(basic_settings, src, dst)

    assert os.path.exists(dst) == True
    assert os.path.exists(src_map) == True

    with io.open(dst, 'r', encoding='utf-8') as f:
        content = f.read()

    with io.open(src_map, 'r', encoding='utf-8') as f:
        sourcemap = json.load(f)

    #print content
    #print "."*60
    #print sourcemap

    # Assert compiled file is ok
    assert content == ("""#content { color: #ff0000; font-weight: bold; }\n\n"""
                      """/*# sourceMappingURL=css/app.map */""")

    # Assert sourcemap is ok
    # Drop 'version' key since it will cause problem with futur libsass
    # versions
    del sourcemap['version']
    # Assert source map is ok
    assert sourcemap == {
        "file": "../app.css",
        "sources": [
                "../app.scss"
        ],
        "mappings": ("AAAA,AAAA,QAAQ,CAAA,EAAE,KAAK,EAAC,OAAO,EAAG,WAAW,EAAC,"
                     "IAAI,GAAK"),
        "names": []
    }

<commit_msg>Fixed map datas in 'tests/100_compiler/002_compile.py::test_002'<commit_after># -*- coding: utf-8 -*-
import os
import io
import json
import pytest

from boussole.conf.model import Settings

def test_001(compiler, temp_builds_dir):
    """compiler.SassCompileHelper.safe_compile: Basic sample without source map"""
    basedir = temp_builds_dir.join('compiler_safecompile_001').strpath

    basic_settings = Settings(initial={
        'SOURCES_PATH': 'scss',
        'TARGET_PATH': 'css',
        'SOURCE_MAP': False,
        'OUTPUT_STYLES': 'compact',
    })

    sourcedir = os.path.join(basedir, basic_settings.SOURCES_PATH)
    targetdir = os.path.join(basedir, basic_settings.TARGET_PATH)

    src = os.path.join(sourcedir, "app.scss")
    dst = os.path.join(targetdir, "app.css")
    src_map = os.path.join(targetdir, "app.map")

    os.makedirs(sourcedir)
    os.makedirs(targetdir)

    # Create sample source to compile
    with io.open(src, 'w', encoding='utf-8') as f:
        result = f.write(
            u"""#content{
            color: #ff0000;
            font-weight: bold;
            &.foo{ border: 1px solid #000000; }
            }
            """)

    # Compile
    success, message = compiler.safe_compile(basic_settings, src, dst)

    assert os.path.exists(dst) == True
    assert os.path.exists(src_map) == False

    # Assert compiled file is ok
    with io.open(dst, 'r', encoding='utf-8') as f:
        content = f.read()

    assert content == ("""#content { color: #ff0000; font-weight: bold; }\n\n"""
                      """#content.foo { border: 1px solid #000000; }\n""")

def test_002(compiler, temp_builds_dir):
    """compiler.SassCompileHelper.safe_compile: Same as basic sample but with
       source map"""
    basedir = temp_builds_dir.join('compiler_safecompile_002').strpath

    basic_settings = Settings(initial={
        'SOURCES_PATH': '.',
        'TARGET_PATH': 'css',
        'SOURCE_MAP': True,
        'OUTPUT_STYLES': 'compact',
    })

    sourcedir = os.path.normpath(os.path.join(basedir, basic_settings.SOURCES_PATH))
    targetdir = os.path.normpath(os.path.join(basedir, basic_settings.TARGET_PATH))

    os.makedirs(sourcedir)
    os.makedirs(targetdir)

    src = os.path.join(sourcedir, "app.scss")
    dst = os.path.join(targetdir, "app.css")
    src_map = os.path.join(targetdir, "app.map")

    # Create sample source to compile
    with io.open(src, 'w', encoding='utf-8') as f:
        result = f.write(u"""#content{ color:#ff0000; font-weight:bold; }""")

    # Compile
    success, message = compiler.safe_compile(basic_settings, src, dst)

    assert os.path.exists(dst) == True
    assert os.path.exists(src_map) == True

    with io.open(dst, 'r', encoding='utf-8') as f:
        content = f.read()

    with io.open(src_map, 'r', encoding='utf-8') as f:
        sourcemap = json.load(f)

    #print content
    #print "."*60
    #print sourcemap

    # Assert compiled file is ok
    assert content == ("""#content { color: #ff0000; font-weight: bold; }\n\n"""
                      """/*# sourceMappingURL=css/app.map */""")

    # Assert sourcemap is ok
    # Drop 'version' key since it will cause problem with futur libsass
    # versions
    del sourcemap['version']
    # Assert source map is ok
    assert sourcemap == {
        "file": "../app.css",
        "sources": [
                "../app.scss"
        ],
        "mappings": ("AAAA,AAAA,QAAQ,CAAA,EAAE,KAAK,EAAC,OAAO,EAAE,WAAW,EAAC,"
                     "IAAI,GAAI"),
        "names": []
    }

<|endoftext|>


<commit_before>import time
import string
import re
import keyword
from Tkinter import *
from Delegator import Delegator
import IdlePrefs

#$ event <<toggle-auto-coloring>>
#$ win <Control-slash>
#$ unix <Control-slash>

__debug__ = 0

def any(name, list):
    return "(?P<%s>" % name + string.join(list, "|") + ")"

def make_pat():
    kw = r"\b" + any("KEYWORD", keyword.kwlist) + r"\b"
    comment = any("COMMENT", [r"#[^\n]*"])
    sqstring = r"(\b[rR])?'[^'\\\n]*(\\.[^'\\\n]*)*'?"
    dqstring = r'(\b[rR])?"[^"\\\n]*(\\.[^"\\\n]*)*"?'
    sq3string = r"(\b[rR])?'''[^'\\]*((\\.|'(?!''))[^'\\]*)*(''')?"
    dq3string = r'(\b[rR])?"""[^"\\]*((\\.|"(?!""))[^"\\]*)*(""")?'
    string = any("STRING", [sq3string, dq3string, sqstring, dqstring])
    return kw + "|" + comment + "|" + string + "|" + any("SYNC", [r"\n"])

prog = re.compile(make_pat(), re.S)
idprog = re.compile(r"\s+(\w+)", re.S)

class ColorDelegator(Delegator):

    def __init__(self):
        Delegator.__init__(self)
        self.prog = prog
        self.idprog = idprog

    def setdelegate(self, delegate):
        if self.delegate is not None:
            self.unbind("<<toggle-auto-coloring>>")
        Delegator.setdelegate(self, delegate)
        if delegate is not None:
            self.config_colors()
            self.bind("<<toggle-auto-coloring>>", self.toggle_colorize_event)
            self.notify_range("1.0", "end")

    def config_colors(self):
        for tag, cnf in self.tagdefs.items():
            if cnf:
                apply(self.tag_configure, (tag,), cnf)
        self.tag_raise('sel')

    cprefs = IdlePrefs.ColorPrefs()

    tagdefs = {
        "COMMENT":    {"foreground": cprefs.CComment[0],
                       "background": cprefs.CComment[1]},
        "KEYWORD":    {"foreground": cprefs.CKeyword[0],
                       "background": cprefs.CKeyword[1]},
        "STRING":     {"foreground": cprefs.CString[0],
                       "background": cprefs.CString[1]},
        "DEFINITION": {"foreground": cprefs.CDefinition[0],
                       "background": cprefs.CDefinition[1]},

        "SYNC":       {"background": cprefs.CSync[0],
                       "background": cprefs.CSync[1]},
        "TODO":       {"background": cprefs.CTodo[0],
                       "background": cprefs.CTodo[1]},

        "BREAK":      {"background": cprefs.CBreak[0],
                       "background": cprefs.CBreak[1]},

        # The following is used by ReplaceDialog:
        "hit":        {"foreground": cprefs.CHit[0],
                       "background": cprefs.CHit[1]},
        }

    def insert(self, index, chars, tags=None):
        index = self.index(index)
        self.delegate.insert(index, chars, tags)
        self.notify_range(index, index + "+%dc" % len(chars))

    def delete(self, index1, index2=None):
        index1 = self.index(index1)
        self.delegate.delete(index1, index2)
        self.notify_range(index1)

    after_id = None
    allow_colorizing = 1
    colorizing = 0

    def notify_range(self, index1, index2=None):
        self.tag_add("TODO", index1, index2)
        if self.after_id:
            if __debug__: print "colorizing already scheduled"
            return
        if self.colorizing:
            self.stop_colorizing = 1
            if __debug__: print "stop colorizing"
        if self.allow_colorizing:
            if __debug__: print "schedule colorizing"
            self.after_id = self.after(1, self.recolorize)

    close_when_done = None # Window to be closed when done colorizing

    def close(self, close_when_done=None):
        if self.after_id:
            after_id = self.after_id
            self.after_id = None
            if __debug__: print "cancel scheduled recolorizer"
            self.after_cancel(after_id)
        self.allow_colorizing = 0
        self.stop_colorizing = 1
        if close_when_done:
            if not self.colorizing:
                close_when_done.destroy()
            else:
                self.close_when_done = close_when_done

    def toggle_colorize_event(self, event):
        if self.after_id:
            after_id = self.after_id
            self.after_id = None
            if __debug__: print "cancel scheduled recolorizer"
            self.after_cancel(after_id)
        if self.allow_colorizing and self.colorizing:
            if __debug__: print "stop colorizing"
            self.stop_colorizing = 1
        self.allow_colorizing = not self.allow_colorizing
        if self.allow_colorizing and not self.colorizing:
            self.after_id = self.after(1, self.recolorize)
        if __debug__:
            print "auto colorizing turned", self.allow_colorizing and "on" or "off"
        return "break"

    def recolorize(self):
        self.after_id = None
        if not self.delegate:
            if __debug__: print "no delegate"
            return
        if not self.allow_colorizing:
            if __debug__: print "auto colorizing is off"
            return
        if self.colorizing:
            if __debug__: print "already colorizing"
            return
        try:
            self.stop_colorizing = 0
            self.colorizing = 1
            if __debug__: print "colorizing..."
            t0 = time.clock()
            self.recolorize_main()
            t1 = time.clock()
            if __debug__: print "%.3f seconds" % (t1-t0)
        finally:
            self.colorizing = 0
        if self.allow_colorizing and self.tag_nextrange("TODO", "1.0"):
            if __debug__: print "reschedule colorizing"
            self.after_id = self.after(1, self.recolorize)
        if self.close_when_done:
            top = self.close_when_done
            self.close_when_done = None
            top.destroy()

    def recolorize_main(self):
        next = "1.0"
        was_ok = is_ok = 0
        while 1:
            item = self.tag_nextrange("TODO", next)
            if not item:
                break
            head, tail = item
            self.tag_remove("SYNC", head, tail)
            item = self.tag_prevrange("SYNC", head)
            if item:
                head = item[1]
            else:
                head = "1.0"

            chars = ""
            mark = head
            is_ok = was_ok = 0
            while not (was_ok and is_ok):
                next = self.index(mark + " lineend +1c")
                was_ok = "SYNC" in self.tag_names(next + "-1c")
                line = self.get(mark, next)
                ##print head, "get", mark, next, "->", `line`
                if not line:
                    return
                for tag in self.tagdefs.keys():
                    self.tag_remove(tag, mark, next)
                chars = chars + line
                m = self.prog.search(chars)
                while m:
                    i, j = m.span()
                    for key, value in m.groupdict().items():
                        if value:
                            a, b = m.span(key)
                            self.tag_add(key,
                                         head + "+%dc" % a,
                                         head + "+%dc" % b)
                            if value in ("def", "class"):
                                m1 = self.idprog.match(chars, b)
                                if m1:
                                    a, b = m1.span(1)
                                    self.tag_add("DEFINITION",
                                                 head + "+%dc" % a,
                                                 head + "+%dc" % b)
                    m = self.prog.search(chars, j)
                is_ok = "SYNC" in self.tag_names(next + "-1c")
                mark = next
                if is_ok:
                    head = mark
                    chars = ""
                self.update()
                if self.stop_colorizing:
                    if __debug__: print "colorizing stopped"
                    return

def main():
    from Percolator import Percolator
    root = Tk()
    root.wm_protocol("WM_DELETE_WINDOW", root.quit)
    text = Text(background="white")
    text.pack(expand=1, fill="both")
    text.focus_set()
    p = Percolator(text)
    d = ColorDelegator()
    p.insertfilter(d)
    root.mainloop()

if __name__ == "__main__":
    main()
<commit_msg>Super-elegant patch by Tim Peters that speeds up colorization dramatically (up to 15 times he claims).  Works by reading more than one line at a time, up to 100-line chunks (starting with one line and then doubling up to the limit).  On a typical machine (e.g. Tim's P5-166) this doesn't reduce interactive responsiveness in a noticeable way.<commit_after>import time
import string
import re
import keyword
from Tkinter import *
from Delegator import Delegator
import IdlePrefs

#$ event <<toggle-auto-coloring>>
#$ win <Control-slash>
#$ unix <Control-slash>

__debug__ = 0

def any(name, list):
    return "(?P<%s>" % name + string.join(list, "|") + ")"

def make_pat():
    kw = r"\b" + any("KEYWORD", keyword.kwlist) + r"\b"
    comment = any("COMMENT", [r"#[^\n]*"])
    sqstring = r"(\b[rR])?'[^'\\\n]*(\\.[^'\\\n]*)*'?"
    dqstring = r'(\b[rR])?"[^"\\\n]*(\\.[^"\\\n]*)*"?'
    sq3string = r"(\b[rR])?'''[^'\\]*((\\.|'(?!''))[^'\\]*)*(''')?"
    dq3string = r'(\b[rR])?"""[^"\\]*((\\.|"(?!""))[^"\\]*)*(""")?'
    string = any("STRING", [sq3string, dq3string, sqstring, dqstring])
    return kw + "|" + comment + "|" + string + "|" + any("SYNC", [r"\n"])

prog = re.compile(make_pat(), re.S)
idprog = re.compile(r"\s+(\w+)", re.S)

class ColorDelegator(Delegator):

    def __init__(self):
        Delegator.__init__(self)
        self.prog = prog
        self.idprog = idprog

    def setdelegate(self, delegate):
        if self.delegate is not None:
            self.unbind("<<toggle-auto-coloring>>")
        Delegator.setdelegate(self, delegate)
        if delegate is not None:
            self.config_colors()
            self.bind("<<toggle-auto-coloring>>", self.toggle_colorize_event)
            self.notify_range("1.0", "end")

    def config_colors(self):
        for tag, cnf in self.tagdefs.items():
            if cnf:
                apply(self.tag_configure, (tag,), cnf)
        self.tag_raise('sel')

    cprefs = IdlePrefs.ColorPrefs()

    tagdefs = {
        "COMMENT":    {"foreground": cprefs.CComment[0],
                       "background": cprefs.CComment[1]},
        "KEYWORD":    {"foreground": cprefs.CKeyword[0],
                       "background": cprefs.CKeyword[1]},
        "STRING":     {"foreground": cprefs.CString[0],
                       "background": cprefs.CString[1]},
        "DEFINITION": {"foreground": cprefs.CDefinition[0],
                       "background": cprefs.CDefinition[1]},

        "SYNC":       {"background": cprefs.CSync[0],
                       "background": cprefs.CSync[1]},
        "TODO":       {"background": cprefs.CTodo[0],
                       "background": cprefs.CTodo[1]},

        "BREAK":      {"background": cprefs.CBreak[0],
                       "background": cprefs.CBreak[1]},

        # The following is used by ReplaceDialog:
        "hit":        {"foreground": cprefs.CHit[0],
                       "background": cprefs.CHit[1]},
        }

    def insert(self, index, chars, tags=None):
        index = self.index(index)
        self.delegate.insert(index, chars, tags)
        self.notify_range(index, index + "+%dc" % len(chars))

    def delete(self, index1, index2=None):
        index1 = self.index(index1)
        self.delegate.delete(index1, index2)
        self.notify_range(index1)

    after_id = None
    allow_colorizing = 1
    colorizing = 0

    def notify_range(self, index1, index2=None):
        self.tag_add("TODO", index1, index2)
        if self.after_id:
            if __debug__: print "colorizing already scheduled"
            return
        if self.colorizing:
            self.stop_colorizing = 1
            if __debug__: print "stop colorizing"
        if self.allow_colorizing:
            if __debug__: print "schedule colorizing"
            self.after_id = self.after(1, self.recolorize)

    close_when_done = None # Window to be closed when done colorizing

    def close(self, close_when_done=None):
        if self.after_id:
            after_id = self.after_id
            self.after_id = None
            if __debug__: print "cancel scheduled recolorizer"
            self.after_cancel(after_id)
        self.allow_colorizing = 0
        self.stop_colorizing = 1
        if close_when_done:
            if not self.colorizing:
                close_when_done.destroy()
            else:
                self.close_when_done = close_when_done

    def toggle_colorize_event(self, event):
        if self.after_id:
            after_id = self.after_id
            self.after_id = None
            if __debug__: print "cancel scheduled recolorizer"
            self.after_cancel(after_id)
        if self.allow_colorizing and self.colorizing:
            if __debug__: print "stop colorizing"
            self.stop_colorizing = 1
        self.allow_colorizing = not self.allow_colorizing
        if self.allow_colorizing and not self.colorizing:
            self.after_id = self.after(1, self.recolorize)
        if __debug__:
            print "auto colorizing turned", self.allow_colorizing and "on" or "off"
        return "break"

    def recolorize(self):
        self.after_id = None
        if not self.delegate:
            if __debug__: print "no delegate"
            return
        if not self.allow_colorizing:
            if __debug__: print "auto colorizing is off"
            return
        if self.colorizing:
            if __debug__: print "already colorizing"
            return
        try:
            self.stop_colorizing = 0
            self.colorizing = 1
            if __debug__: print "colorizing..."
            t0 = time.clock()
            self.recolorize_main()
            t1 = time.clock()
            if __debug__: print "%.3f seconds" % (t1-t0)
        finally:
            self.colorizing = 0
        if self.allow_colorizing and self.tag_nextrange("TODO", "1.0"):
            if __debug__: print "reschedule colorizing"
            self.after_id = self.after(1, self.recolorize)
        if self.close_when_done:
            top = self.close_when_done
            self.close_when_done = None
            top.destroy()

    def recolorize_main(self):
        next = "1.0"
        was_ok = is_ok = 0
        while 1:
            item = self.tag_nextrange("TODO", next)
            if not item:
                break
            head, tail = item
            self.tag_remove("SYNC", head, tail)
            item = self.tag_prevrange("SYNC", head)
            if item:
                head = item[1]
            else:
                head = "1.0"

            chars = ""
            mark = head
            lines_to_get = 1
            is_ok = was_ok = 0
            while not (was_ok and is_ok):
                next = self.index(mark + "+%d lines linestart" %
                                         lines_to_get)
                lines_to_get = min(lines_to_get * 2, 100)
                was_ok = "SYNC" in self.tag_names(next + "-1c")
                line = self.get(mark, next)
                ##print head, "get", mark, next, "->", `line`
                if not line:
                    return
                for tag in self.tagdefs.keys():
                    self.tag_remove(tag, mark, next)
                chars = chars + line
                m = self.prog.search(chars)
                while m:
                    i, j = m.span()
                    for key, value in m.groupdict().items():
                        if value:
                            a, b = m.span(key)
                            self.tag_add(key,
                                         head + "+%dc" % a,
                                         head + "+%dc" % b)
                            if value in ("def", "class"):
                                m1 = self.idprog.match(chars, b)
                                if m1:
                                    a, b = m1.span(1)
                                    self.tag_add("DEFINITION",
                                                 head + "+%dc" % a,
                                                 head + "+%dc" % b)
                    m = self.prog.search(chars, j)
                is_ok = "SYNC" in self.tag_names(next + "-1c")
                mark = next
                if is_ok:
                    head = mark
                    chars = ""
                self.update()
                if self.stop_colorizing:
                    if __debug__: print "colorizing stopped"
                    return

def main():
    from Percolator import Percolator
    root = Tk()
    root.wm_protocol("WM_DELETE_WINDOW", root.quit)
    text = Text(background="white")
    text.pack(expand=1, fill="both")
    text.focus_set()
    p = Percolator(text)
    d = ColorDelegator()
    p.insertfilter(d)
    root.mainloop()

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>import sys
import re
global default_var
default_var = {}

class ParameterGroup(object):
    """
    Single parameter group
    """
    def __init__(self, name):
        self.name = name
        self.params = []

    def AddParameter(self, param):
        """
        Add parameter to the group
        """
        self.params.append(param)

    def GetName(self):
        """
        Get parameter group name
        """
        return self.name

    def GetParams(self):
        """
        Returns the parsed list of parameters. Every parameter is a Parameter
        object. Note that returned object is not a copy. Modifications affect
        state of the parser.
        """
        return sorted(self.params,
                key=lambda x: x.GetFieldValue("code"))

class Parameter(object):
    """
    Single parameter
    """

    # Define sorting order of the fields
    priority = {
		"board": 9,
        "short_desc": 8,
        "long_desc": 7,
        "min": 5,
        "max": 4,
        "unit": 3,
        "decimal": 2,
        # all others == 0 (sorted alphabetically)
    }

    def __init__(self, name, type, default = ""):
        self.fields = {}
        self.values = {}
        self.name = name
        self.type = type
        self.default = default

    def GetName(self):
        return self.name

    def GetType(self):
        return self.type

    def GetDefault(self):
        return self.default

    def SetField(self, code, value):
        """
        Set named field value
        """
        self.fields[code] = value

    def SetEnumValue(self, code, value):
        """
        Set named enum value
        """
        self.values[code] = value

    def GetFieldCodes(self):
        """
        Return list of existing field codes in convenient order
        """
        keys = self.fields.keys()
        keys = sorted(keys)
        keys = sorted(keys, key=lambda x: self.priority.get(x, 0), reverse=True)
        return keys

    def GetFieldValue(self, code):
        """
        Return value of the given field code or None if not found.
        """
        fv =  self.fields.get(code)
        if not fv:
                # required because python 3 sorted does not accept None
                return ""
        return fv

    def GetEnumCodes(self):
        """
        Return list of existing value codes in convenient order
        """
        keys = self.values.keys()
        #keys = sorted(keys)
        #keys = sorted(keys, key=lambda x: self.priority.get(x, 0), reverse=True)
        return keys

    def GetEnumValue(self, code):
        """
        Return value of the given enum code or None if not found.
        """
        fv =  self.values.get(code)
        if not fv:
                # required because python 3 sorted does not accept None
                return ""
        return fv

class SourceParser(object):
    """
    Parses provided data and stores all found parameters internally.
    """

    re_split_lines = re.compile(r'[\r\n]+')
    re_comment_start = re.compile(r'^\/\*\*')
    re_comment_content = re.compile(r'^\*\s*(.*)')
    re_comment_tag = re.compile(r'@([a-zA-Z][a-zA-Z0-9_]*)\s*(.*)')
    re_comment_end = re.compile(r'(.*?)\s*\*\/')
    re_parameter_definition = re.compile(r'PARAM_DEFINE_([A-Z_][A-Z0-9_]*)\s*\(([A-Z_][A-Z0-9_]*)\s*,\s*([^ ,\)]+)\s*\)\s*;')
    re_px4_parameter_definition = re.compile(r'PX4_PARAM_DEFINE_([A-Z_][A-Z0-9_]*)\s*\(([A-Z_][A-Z0-9_]*)\s*\)\s*;')
    re_px4_param_default_definition = re.compile(r'#define\s*PARAM_([A-Z_][A-Z0-9_]*)\s*([^ ,\)]+)\s*')
    re_cut_type_specifier = re.compile(r'[a-z]+$')
    re_is_a_number = re.compile(r'^-?[0-9\.]')
    re_remove_dots = re.compile(r'\.+$')
    re_remove_carriage_return = re.compile('\n+')

    valid_tags = set(["group", "board", "min", "max", "unit", "decimal", "increment", "reboot_required", "value"])

    # Order of parameter groups
    priority = {
        # All other groups = 0 (sort alphabetically)
        "Miscellaneous": -10
    }

    def __init__(self):
        self.param_groups = {}

    def Parse(self, contents):
        """
        Incrementally parse program contents and append all found parameters
        to the list.
        """
        # This code is essentially a comment-parsing grammar. "state"
        # represents parser state. It contains human-readable state
        # names.
        state = None
        for line in self.re_split_lines.split(contents):
            line = line.strip()
            # Ignore empty lines
            if line == "":
                continue
            if self.re_comment_start.match(line):
                state = "wait-short"
                short_desc = None
                long_desc = None
                tags = {}
                def_values = {}
            elif state is not None and state != "comment-processed":
                m = self.re_comment_end.search(line)
                if m:
                    line = m.group(1)
                    last_comment_line = True
                else:
                    last_comment_line = False
                m = self.re_comment_content.match(line)
                if m:
                    comment_content = m.group(1)
                    if comment_content == "":
                        # When short comment ends with empty comment line,
                        # start waiting for the next part - long comment.
                        if state == "wait-short-end":
                            state = "wait-long"
                    else:
                        m = self.re_comment_tag.match(comment_content)
                        if m:
                            tag, desc = m.group(1, 2)
                            if (tag == "value"):
                                # Take the meta info string and split the code and description
                                metainfo = desc.split(" ",  1)
                                def_values[metainfo[0]] = metainfo[1]
                            else:
                                tags[tag] = desc
                            current_tag = tag
                            state = "wait-tag-end"
                        elif state == "wait-short":
                            # Store first line of the short description
                            short_desc = comment_content
                            state = "wait-short-end"
                        elif state == "wait-short-end":
                            # Append comment line to the short description
                            short_desc += "\n" + comment_content
                        elif state == "wait-long":
                            # Store first line of the long description
                            long_desc = comment_content
                            state = "wait-long-end"
                        elif state == "wait-long-end":
                            # Append comment line to the long description
                            long_desc += "\n" + comment_content
                        elif state == "wait-tag-end":
                            # Append comment line to the tag text
                            tags[current_tag] += "\n" + comment_content
                        else:
                            raise AssertionError(
                                    "Invalid parser state: %s" % state)
                elif not last_comment_line:
                    # Invalid comment line (inside comment, but not starting with
                    # "*" or "*/". Reset parsed content.
                    state = None
                if last_comment_line:
                    state = "comment-processed"
            else:
                tp = None
                name = None
                defval = ""
                # Non-empty line outside the comment
                m = self.re_px4_param_default_definition.match(line)
                # Default value handling
                if m:
                    name_m, defval_m = m.group(1,2)
                    default_var[name_m] = defval_m
                m = self.re_parameter_definition.match(line)
                if m:
                    tp, name, defval = m.group(1, 2, 3)
                else:
                    m = self.re_px4_parameter_definition.match(line)
                    if m:
                        tp, name = m.group(1, 2)
                        if (name+'_DEFAULT') in default_var:
                            defval = default_var[name+'_DEFAULT']
                if tp is not None:
                    # Remove trailing type specifier from numbers: 0.1f => 0.1
                    if defval != "" and self.re_is_a_number.match(defval):
                        defval = self.re_cut_type_specifier.sub('', defval)
                    param = Parameter(name, tp, defval)
                    param.SetField("short_desc", name)
                    # If comment was found before the parameter declaration,
                    # inject its data into the newly created parameter.
                    group = "Miscellaneous"
                    if state == "comment-processed":
                        if short_desc is not None:
                            param.SetField("short_desc", self.re_remove_dots.sub('', short_desc))
                        if long_desc is not None:
                            long_desc = self.re_remove_carriage_return.sub(' ', long_desc)
                            param.SetField("long_desc", long_desc)
                        for tag in tags:
                            if tag == "group":
                                group = tags[tag]
                            elif tag not in self.valid_tags:
                                sys.stderr.write("Skipping invalid documentation tag: '%s'\n" % tag)
                                return False
                            else:
                                param.SetField(tag, tags[tag])
                        for def_value in def_values:
                            param.SetEnumValue(def_value, def_values[def_value])
                    # Store the parameter
                    if group not in self.param_groups:
                        self.param_groups[group] = ParameterGroup(group)
                    self.param_groups[group].AddParameter(param)
                state = None
        return True

    def IsNumber(self, numberString):
        try:
            float(numberString)
            return True
        except ValueError:
            return False

    def Validate(self):
        """
        Validates the parameter meta data.
        """
        seenParamNames = []
        for group in self.GetParamGroups():
            for param in group.GetParams():
                name  = param.GetName()
                board = param.GetFieldValue("board")
                # Check for duplicates
                name_plus_board = name + "+" + board
                for seenParamName in seenParamNames:
                    if seenParamName == name_plus_board:
                        sys.stderr.write("Duplicate parameter definition: {0}\n".format(name_plus_board))
                        return False
                seenParamNames.append(name_plus_board)
                # Validate values
                default = param.GetDefault()
                min = param.GetFieldValue("min")
                max = param.GetFieldValue("max")
                #sys.stderr.write("{0} default:{1} min:{2} max:{3}\n".format(name, default, min, max))
                if default != "" and not self.IsNumber(default):
                    sys.stderr.write("Default value not number: {0} {1}\n".format(name, default))
                    return False
                if min != "":
                    if not self.IsNumber(min):
                        sys.stderr.write("Min value not number: {0} {1}\n".format(name, min))
                        return False
                    if default != "" and float(default) < float(min):
                        sys.stderr.write("Default value is smaller than min: {0} default:{1} min:{2}\n".format(name, default, min))
                        return False
                if max != "":
                    if not self.IsNumber(max):
                        sys.stderr.write("Max value not number: {0} {1}\n".format(name, max))
                        return False
                    if default != "" and float(default) > float(max):
                        sys.stderr.write("Default value is larger than max: {0} default:{1} max:{2}\n".format(name, default, max))
                        return False
                for code in param.GetEnumCodes():
                        if not self.IsNumber(code):
                            sys.stderr.write("Min value not number: {0} {1}\n".format(name, code))
                            return False
                        if param.GetEnumValue(code) == "":
                            sys.stderr.write("Description for enum value is empty: {0} {1}\n".format(name, code))
                            return False
        return True

    def GetParamGroups(self):
        """
        Returns the parsed list of parameters. Every parameter is a Parameter
        object. Note that returned object is not a copy. Modifications affect
        state of the parser.
        """
        groups = self.param_groups.values()
        groups = sorted(groups, key=lambda x: x.GetName())
        groups = sorted(groups, key=lambda x: self.priority.get(x.GetName(), 0), reverse=True)
        return groups
<commit_msg>allow boolean tag<commit_after>import sys
import re
global default_var
default_var = {}

class ParameterGroup(object):
    """
    Single parameter group
    """
    def __init__(self, name):
        self.name = name
        self.params = []

    def AddParameter(self, param):
        """
        Add parameter to the group
        """
        self.params.append(param)

    def GetName(self):
        """
        Get parameter group name
        """
        return self.name

    def GetParams(self):
        """
        Returns the parsed list of parameters. Every parameter is a Parameter
        object. Note that returned object is not a copy. Modifications affect
        state of the parser.
        """
        return sorted(self.params,
                key=lambda x: x.GetFieldValue("code"))

class Parameter(object):
    """
    Single parameter
    """

    # Define sorting order of the fields
    priority = {
		"board": 9,
        "short_desc": 8,
        "long_desc": 7,
        "min": 5,
        "max": 4,
        "unit": 3,
        "decimal": 2,
        # all others == 0 (sorted alphabetically)
    }

    def __init__(self, name, type, default = ""):
        self.fields = {}
        self.values = {}
        self.name = name
        self.type = type
        self.default = default

    def GetName(self):
        return self.name

    def GetType(self):
        return self.type

    def GetDefault(self):
        return self.default

    def SetField(self, code, value):
        """
        Set named field value
        """
        self.fields[code] = value

    def SetEnumValue(self, code, value):
        """
        Set named enum value
        """
        self.values[code] = value

    def GetFieldCodes(self):
        """
        Return list of existing field codes in convenient order
        """
        keys = self.fields.keys()
        keys = sorted(keys)
        keys = sorted(keys, key=lambda x: self.priority.get(x, 0), reverse=True)
        return keys

    def GetFieldValue(self, code):
        """
        Return value of the given field code or None if not found.
        """
        fv =  self.fields.get(code)
        if not fv:
                # required because python 3 sorted does not accept None
                return ""
        return fv

    def GetEnumCodes(self):
        """
        Return list of existing value codes in convenient order
        """
        keys = self.values.keys()
        #keys = sorted(keys)
        #keys = sorted(keys, key=lambda x: self.priority.get(x, 0), reverse=True)
        return keys

    def GetEnumValue(self, code):
        """
        Return value of the given enum code or None if not found.
        """
        fv =  self.values.get(code)
        if not fv:
                # required because python 3 sorted does not accept None
                return ""
        return fv

class SourceParser(object):
    """
    Parses provided data and stores all found parameters internally.
    """

    re_split_lines = re.compile(r'[\r\n]+')
    re_comment_start = re.compile(r'^\/\*\*')
    re_comment_content = re.compile(r'^\*\s*(.*)')
    re_comment_tag = re.compile(r'@([a-zA-Z][a-zA-Z0-9_]*)\s*(.*)')
    re_comment_end = re.compile(r'(.*?)\s*\*\/')
    re_parameter_definition = re.compile(r'PARAM_DEFINE_([A-Z_][A-Z0-9_]*)\s*\(([A-Z_][A-Z0-9_]*)\s*,\s*([^ ,\)]+)\s*\)\s*;')
    re_px4_parameter_definition = re.compile(r'PX4_PARAM_DEFINE_([A-Z_][A-Z0-9_]*)\s*\(([A-Z_][A-Z0-9_]*)\s*\)\s*;')
    re_px4_param_default_definition = re.compile(r'#define\s*PARAM_([A-Z_][A-Z0-9_]*)\s*([^ ,\)]+)\s*')
    re_cut_type_specifier = re.compile(r'[a-z]+$')
    re_is_a_number = re.compile(r'^-?[0-9\.]')
    re_remove_dots = re.compile(r'\.+$')
    re_remove_carriage_return = re.compile('\n+')

    valid_tags = set(["group", "board", "min", "max", "unit", "decimal", "increment", "reboot_required", "value", "boolean"])

    # Order of parameter groups
    priority = {
        # All other groups = 0 (sort alphabetically)
        "Miscellaneous": -10
    }

    def __init__(self):
        self.param_groups = {}

    def Parse(self, contents):
        """
        Incrementally parse program contents and append all found parameters
        to the list.
        """
        # This code is essentially a comment-parsing grammar. "state"
        # represents parser state. It contains human-readable state
        # names.
        state = None
        for line in self.re_split_lines.split(contents):
            line = line.strip()
            # Ignore empty lines
            if line == "":
                continue
            if self.re_comment_start.match(line):
                state = "wait-short"
                short_desc = None
                long_desc = None
                tags = {}
                def_values = {}
            elif state is not None and state != "comment-processed":
                m = self.re_comment_end.search(line)
                if m:
                    line = m.group(1)
                    last_comment_line = True
                else:
                    last_comment_line = False
                m = self.re_comment_content.match(line)
                if m:
                    comment_content = m.group(1)
                    if comment_content == "":
                        # When short comment ends with empty comment line,
                        # start waiting for the next part - long comment.
                        if state == "wait-short-end":
                            state = "wait-long"
                    else:
                        m = self.re_comment_tag.match(comment_content)
                        if m:
                            tag, desc = m.group(1, 2)
                            if (tag == "value"):
                                # Take the meta info string and split the code and description
                                metainfo = desc.split(" ",  1)
                                def_values[metainfo[0]] = metainfo[1]
                            else:
                                tags[tag] = desc
                            current_tag = tag
                            state = "wait-tag-end"
                        elif state == "wait-short":
                            # Store first line of the short description
                            short_desc = comment_content
                            state = "wait-short-end"
                        elif state == "wait-short-end":
                            # Append comment line to the short description
                            short_desc += "\n" + comment_content
                        elif state == "wait-long":
                            # Store first line of the long description
                            long_desc = comment_content
                            state = "wait-long-end"
                        elif state == "wait-long-end":
                            # Append comment line to the long description
                            long_desc += "\n" + comment_content
                        elif state == "wait-tag-end":
                            # Append comment line to the tag text
                            tags[current_tag] += "\n" + comment_content
                        else:
                            raise AssertionError(
                                    "Invalid parser state: %s" % state)
                elif not last_comment_line:
                    # Invalid comment line (inside comment, but not starting with
                    # "*" or "*/". Reset parsed content.
                    state = None
                if last_comment_line:
                    state = "comment-processed"
            else:
                tp = None
                name = None
                defval = ""
                # Non-empty line outside the comment
                m = self.re_px4_param_default_definition.match(line)
                # Default value handling
                if m:
                    name_m, defval_m = m.group(1,2)
                    default_var[name_m] = defval_m
                m = self.re_parameter_definition.match(line)
                if m:
                    tp, name, defval = m.group(1, 2, 3)
                else:
                    m = self.re_px4_parameter_definition.match(line)
                    if m:
                        tp, name = m.group(1, 2)
                        if (name+'_DEFAULT') in default_var:
                            defval = default_var[name+'_DEFAULT']
                if tp is not None:
                    # Remove trailing type specifier from numbers: 0.1f => 0.1
                    if defval != "" and self.re_is_a_number.match(defval):
                        defval = self.re_cut_type_specifier.sub('', defval)
                    param = Parameter(name, tp, defval)
                    param.SetField("short_desc", name)
                    # If comment was found before the parameter declaration,
                    # inject its data into the newly created parameter.
                    group = "Miscellaneous"
                    if state == "comment-processed":
                        if short_desc is not None:
                            param.SetField("short_desc", self.re_remove_dots.sub('', short_desc))
                        if long_desc is not None:
                            long_desc = self.re_remove_carriage_return.sub(' ', long_desc)
                            param.SetField("long_desc", long_desc)
                        for tag in tags:
                            if tag == "group":
                                group = tags[tag]
                            elif tag not in self.valid_tags:
                                sys.stderr.write("Skipping invalid documentation tag: '%s'\n" % tag)
                                return False
                            else:
                                param.SetField(tag, tags[tag])
                        for def_value in def_values:
                            param.SetEnumValue(def_value, def_values[def_value])
                    # Store the parameter
                    if group not in self.param_groups:
                        self.param_groups[group] = ParameterGroup(group)
                    self.param_groups[group].AddParameter(param)
                state = None
        return True

    def IsNumber(self, numberString):
        try:
            float(numberString)
            return True
        except ValueError:
            return False

    def Validate(self):
        """
        Validates the parameter meta data.
        """
        seenParamNames = []
        for group in self.GetParamGroups():
            for param in group.GetParams():
                name  = param.GetName()
                board = param.GetFieldValue("board")
                # Check for duplicates
                name_plus_board = name + "+" + board
                for seenParamName in seenParamNames:
                    if seenParamName == name_plus_board:
                        sys.stderr.write("Duplicate parameter definition: {0}\n".format(name_plus_board))
                        return False
                seenParamNames.append(name_plus_board)
                # Validate values
                default = param.GetDefault()
                min = param.GetFieldValue("min")
                max = param.GetFieldValue("max")
                #sys.stderr.write("{0} default:{1} min:{2} max:{3}\n".format(name, default, min, max))
                if default != "" and not self.IsNumber(default):
                    sys.stderr.write("Default value not number: {0} {1}\n".format(name, default))
                    return False
                if min != "":
                    if not self.IsNumber(min):
                        sys.stderr.write("Min value not number: {0} {1}\n".format(name, min))
                        return False
                    if default != "" and float(default) < float(min):
                        sys.stderr.write("Default value is smaller than min: {0} default:{1} min:{2}\n".format(name, default, min))
                        return False
                if max != "":
                    if not self.IsNumber(max):
                        sys.stderr.write("Max value not number: {0} {1}\n".format(name, max))
                        return False
                    if default != "" and float(default) > float(max):
                        sys.stderr.write("Default value is larger than max: {0} default:{1} max:{2}\n".format(name, default, max))
                        return False
                for code in param.GetEnumCodes():
                        if not self.IsNumber(code):
                            sys.stderr.write("Min value not number: {0} {1}\n".format(name, code))
                            return False
                        if param.GetEnumValue(code) == "":
                            sys.stderr.write("Description for enum value is empty: {0} {1}\n".format(name, code))
                            return False
        return True

    def GetParamGroups(self):
        """
        Returns the parsed list of parameters. Every parameter is a Parameter
        object. Note that returned object is not a copy. Modifications affect
        state of the parser.
        """
        groups = self.param_groups.values()
        groups = sorted(groups, key=lambda x: x.GetName())
        groups = sorted(groups, key=lambda x: self.priority.get(x.GetName(), 0), reverse=True)
        return groups
<|endoftext|>


<commit_before>#!/usr/bin/python
#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#
#-------------------------------------------------------------

from py4j.protocol import Py4JJavaError, Py4JError
import traceback
import os
from pyspark.sql import DataFrame, SQLContext
from pyspark.rdd import RDD

class MLContext(object):

    """
    Simple wrapper class for MLContext in SystemML.jar
    ...
    Attributes
    ----------
    ml : MLContext
        A reference to the java MLContext
    sc : SparkContext
        The SparkContext that has been specified during initialization
    """

    def __init__(self, sc, *args):
        """
        If initialized with a SparkContext, will connect to the Java MLContext
        class.
        args:
            sc (SparkContext): the current SparkContext
            monitor (boolean=False): Whether to monitor the performance
            forceSpark (boolean=False): Whether to force execution on spark
        returns:
            MLContext: Instance of MLContext
        """

        try:
            monitorPerformance = (args[0] if len(args) > 0 else False)
            setForcedSparkExecType = (args[1] if len(args) > 1 else False)
            self.sc = sc
            self.ml = sc._jvm.org.apache.sysml.api.MLContext(sc._jsc, monitorPerformance, setForcedSparkExecType)
        except Py4JError:
            traceback.print_exc()

    def reset(self):
        """
        Call this method of you want to clear any RDDs set via
        registerInput or registerOutput
        """
        try:
            self.ml.reset()
        except Py4JJavaError:
            traceback.print_exc()

    def execute(self, dmlScriptFilePath, *args):
        """
        Executes the script in spark-mode by passing the arguments to the
        MLContext java class.
        Returns:
            MLOutput: an instance of the MLOutput-class
        """
        numArgs = len(args) + 1
        try:
            if numArgs == 1:
                jmlOut = self.ml.execute(dmlScriptFilePath)
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            elif numArgs == 2:
                jmlOut = self.ml.execute(dmlScriptFilePath, args[0])
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            elif numArgs == 3:
                jmlOut = self.ml.execute(dmlScriptFilePath, args[0], args[1])
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            elif numArgs == 4:
                jmlOut = self.ml.execute(dmlScriptFilePath, args[0], args[1], args[2])
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            else:
                raise TypeError('Arguments do not match MLContext-API')
        except Py4JJavaError:
            traceback.print_exc()

    def executeScript(self, dmlScript, nargs=None, outputs=None, isPyDML=False, configFilePath=None):
        """
        Executes the script in spark-mode by passing the arguments to the
        MLContext java class.
        Returns:
            MLOutput: an instance of the MLOutput-class
        """
        try:
            # Register inputs as needed
            if nargs is not None:
                for key, value in nargs.items():
                    if isinstance(value, DataFrame):
                        self.registerInput(key, value)
                        del nargs[key]
                    else:
                        nargs[key] = str(value)
            else:
                nargs = {}

            # Register outputs as needed
            if outputs is not None:
                for out in outputs:
                    self.registerOutput(out)

            # Execute script
            jml_out = self.ml.executeScript(dmlScript, nargs, isPyDML, configFilePath)
            ml_out = MLOutput(jml_out, self.sc)
            return ml_out
        except Py4JJavaError:
            traceback.print_exc()

    def registerInput(self, varName, src, *args):
        """
        Method to register inputs used by the DML script.
        Supported format:
        1. DataFrame
        2. CSV/Text (as JavaRDD<String> or JavaPairRDD<LongWritable, Text>)
        3. Binary blocked RDD (JavaPairRDD<MatrixIndexes,MatrixBlock>))
        Also overloaded to support metadata information such as format, rlen, clen, ...
        Please note the variable names given below in quotes correspond to the variables in DML script.
        These variables need to have corresponding read/write associated in DML script.
        Currently, only matrix variables are supported through registerInput/registerOutput interface.
        To pass scalar variables, use named/positional arguments (described later) or wrap them into matrix variable.
        """
        numArgs = len(args) + 2

        if hasattr(src, '_jdf'):
            rdd = src._jdf
        elif hasattr(src, '_jrdd'):
            rdd = src._jrdd
        else:
            rdd = src

        try:
            if numArgs == 2:
                self.ml.registerInput(varName, rdd)
            elif numArgs == 3:
                self.ml.registerInput(varName, rdd, args[0])
            elif numArgs == 4:
                self.ml.registerInput(varName, rdd, args[0], args[1])
            elif numArgs == 5:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2])
            elif numArgs == 6:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2], args[3])
            elif numArgs == 7:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2], args[3], args[4])
            elif numArgs == 10:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7])
            else:
                raise TypeError('Arguments do not match MLContext-API')
        except Py4JJavaError:

            traceback.print_exc()

    def registerOutput(self, varName):
        """
        Register output variables used in the DML script
        args:
            varName: (String) The name used in the DML script
        """

        try:
            self.ml.registerOutput(varName)
        except Py4JJavaError:
            traceback.print_exc()

    def getDmlJson(self):
        try:
            return self.ml.getMonitoringUtil().getRuntimeInfoInJSONFormat()
        except Py4JJavaError:
            traceback.print_exc()

class MLOutput(object):

    """
    This is a simple wrapper object that returns the output of execute from MLContext
    ...
    Attributes
    ----------
    jmlOut MLContext:
        A reference to the MLOutput object through py4j
    """

    def __init__(self, jmlOut, sc):
        self.jmlOut = jmlOut
        self.sc = sc

    def getBinaryBlockedRDD(self, varName):
        raise Exception('Not supported in Python MLContext')
		#try:
        #    rdd = RDD(self.jmlOut.getBinaryBlockedRDD(varName), self.sc)
        #    return rdd
        #except Py4JJavaError:
        #    traceback.print_exc()

    def getMatrixCharacteristics(self, varName):
        raise Exception('Not supported in Python MLContext')
		#try:
        #    chars = self.jmlOut.getMatrixCharacteristics(varName)
        #    return chars
        #except Py4JJavaError:
        #    traceback.print_exc()

    def getDF(self, sqlContext, varName):
        try:
            jdf = self.jmlOut.getDF(sqlContext._scala_SQLContext, varName)
            df = DataFrame(jdf, sqlContext)
            return df
        except Py4JJavaError:
            traceback.print_exc()

    def getMLMatrix(self, sqlContext, varName):
        raise Exception('Not supported in Python MLContext')
		#try:
        #    mlm = self.jmlOut.getMLMatrix(sqlContext._scala_SQLContext, varName)
        #    return mlm
        #except Py4JJavaError:
        #    traceback.print_exc()

    def getStringRDD(self, varName, format):
		raise Exception('Not supported in Python MLContext')
        #try:
        #    rdd = RDD(self.jmlOut.getStringRDD(varName, format), self.sc)
        #    return rdd
        #except Py4JJavaError:
        #    traceback.print_exc()
<commit_msg>[SYSTEMML-668] Python MLOutput.getDF() Can't Access JVM SQLContext<commit_after>#!/usr/bin/python
#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#
#-------------------------------------------------------------

from py4j.protocol import Py4JJavaError, Py4JError
import traceback
import os
from pyspark.sql import DataFrame, SQLContext
from pyspark.rdd import RDD

class MLContext(object):

    """
    Simple wrapper class for MLContext in SystemML.jar
    ...
    Attributes
    ----------
    ml : MLContext
        A reference to the java MLContext
    sc : SparkContext
        The SparkContext that has been specified during initialization
    """

    def __init__(self, sc, *args):
        """
        If initialized with a SparkContext, will connect to the Java MLContext
        class.
        args:
            sc (SparkContext): the current SparkContext
            monitor (boolean=False): Whether to monitor the performance
            forceSpark (boolean=False): Whether to force execution on spark
        returns:
            MLContext: Instance of MLContext
        """

        try:
            monitorPerformance = (args[0] if len(args) > 0 else False)
            setForcedSparkExecType = (args[1] if len(args) > 1 else False)
            self.sc = sc
            self.ml = sc._jvm.org.apache.sysml.api.MLContext(sc._jsc, monitorPerformance, setForcedSparkExecType)
        except Py4JError:
            traceback.print_exc()

    def reset(self):
        """
        Call this method of you want to clear any RDDs set via
        registerInput or registerOutput
        """
        try:
            self.ml.reset()
        except Py4JJavaError:
            traceback.print_exc()

    def execute(self, dmlScriptFilePath, *args):
        """
        Executes the script in spark-mode by passing the arguments to the
        MLContext java class.
        Returns:
            MLOutput: an instance of the MLOutput-class
        """
        numArgs = len(args) + 1
        try:
            if numArgs == 1:
                jmlOut = self.ml.execute(dmlScriptFilePath)
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            elif numArgs == 2:
                jmlOut = self.ml.execute(dmlScriptFilePath, args[0])
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            elif numArgs == 3:
                jmlOut = self.ml.execute(dmlScriptFilePath, args[0], args[1])
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            elif numArgs == 4:
                jmlOut = self.ml.execute(dmlScriptFilePath, args[0], args[1], args[2])
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            else:
                raise TypeError('Arguments do not match MLContext-API')
        except Py4JJavaError:
            traceback.print_exc()

    def executeScript(self, dmlScript, nargs=None, outputs=None, isPyDML=False, configFilePath=None):
        """
        Executes the script in spark-mode by passing the arguments to the
        MLContext java class.
        Returns:
            MLOutput: an instance of the MLOutput-class
        """
        try:
            # Register inputs as needed
            if nargs is not None:
                for key, value in nargs.items():
                    if isinstance(value, DataFrame):
                        self.registerInput(key, value)
                        del nargs[key]
                    else:
                        nargs[key] = str(value)
            else:
                nargs = {}

            # Register outputs as needed
            if outputs is not None:
                for out in outputs:
                    self.registerOutput(out)

            # Execute script
            jml_out = self.ml.executeScript(dmlScript, nargs, isPyDML, configFilePath)
            ml_out = MLOutput(jml_out, self.sc)
            return ml_out
        except Py4JJavaError:
            traceback.print_exc()

    def registerInput(self, varName, src, *args):
        """
        Method to register inputs used by the DML script.
        Supported format:
        1. DataFrame
        2. CSV/Text (as JavaRDD<String> or JavaPairRDD<LongWritable, Text>)
        3. Binary blocked RDD (JavaPairRDD<MatrixIndexes,MatrixBlock>))
        Also overloaded to support metadata information such as format, rlen, clen, ...
        Please note the variable names given below in quotes correspond to the variables in DML script.
        These variables need to have corresponding read/write associated in DML script.
        Currently, only matrix variables are supported through registerInput/registerOutput interface.
        To pass scalar variables, use named/positional arguments (described later) or wrap them into matrix variable.
        """
        numArgs = len(args) + 2

        if hasattr(src, '_jdf'):
            rdd = src._jdf
        elif hasattr(src, '_jrdd'):
            rdd = src._jrdd
        else:
            rdd = src

        try:
            if numArgs == 2:
                self.ml.registerInput(varName, rdd)
            elif numArgs == 3:
                self.ml.registerInput(varName, rdd, args[0])
            elif numArgs == 4:
                self.ml.registerInput(varName, rdd, args[0], args[1])
            elif numArgs == 5:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2])
            elif numArgs == 6:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2], args[3])
            elif numArgs == 7:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2], args[3], args[4])
            elif numArgs == 10:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7])
            else:
                raise TypeError('Arguments do not match MLContext-API')
        except Py4JJavaError:

            traceback.print_exc()

    def registerOutput(self, varName):
        """
        Register output variables used in the DML script
        args:
            varName: (String) The name used in the DML script
        """

        try:
            self.ml.registerOutput(varName)
        except Py4JJavaError:
            traceback.print_exc()

    def getDmlJson(self):
        try:
            return self.ml.getMonitoringUtil().getRuntimeInfoInJSONFormat()
        except Py4JJavaError:
            traceback.print_exc()

class MLOutput(object):

    """
    This is a simple wrapper object that returns the output of execute from MLContext
    ...
    Attributes
    ----------
    jmlOut MLContext:
        A reference to the MLOutput object through py4j
    """

    def __init__(self, jmlOut, sc):
        self.jmlOut = jmlOut
        self.sc = sc

    def getBinaryBlockedRDD(self, varName):
        raise Exception('Not supported in Python MLContext')
		#try:
        #    rdd = RDD(self.jmlOut.getBinaryBlockedRDD(varName), self.sc)
        #    return rdd
        #except Py4JJavaError:
        #    traceback.print_exc()

    def getMatrixCharacteristics(self, varName):
        raise Exception('Not supported in Python MLContext')
		#try:
        #    chars = self.jmlOut.getMatrixCharacteristics(varName)
        #    return chars
        #except Py4JJavaError:
        #    traceback.print_exc()

    def getDF(self, sqlContext, varName):
        try:
            jdf = self.jmlOut.getDF(sqlContext._ssql_ctx, varName)
            df = DataFrame(jdf, sqlContext)
            return df
        except Py4JJavaError:
            traceback.print_exc()

    def getMLMatrix(self, sqlContext, varName):
        raise Exception('Not supported in Python MLContext')
		#try:
        #    mlm = self.jmlOut.getMLMatrix(sqlContext._scala_SQLContext, varName)
        #    return mlm
        #except Py4JJavaError:
        #    traceback.print_exc()

    def getStringRDD(self, varName, format):
		raise Exception('Not supported in Python MLContext')
        #try:
        #    rdd = RDD(self.jmlOut.getStringRDD(varName, format), self.sc)
        #    return rdd
        #except Py4JJavaError:
        #    traceback.print_exc()
<|endoftext|>