<commit_before>def rad_loss_rate(goeslc):
    """
    Calculates and adds solar radiative loss rate to a GOESLightCurve.

    Extended Summary
    ----------------
    This function calculates the radiative loss rate as a function of
    time of solar coronal soft X-ray-emitting plasma across all
    wavelengths given a GOESLightCurve object.  The units of the
    results are erg/s. This is done by calling calc_rad_loss().
    For more information see documentation in that function.  Once
    the radiative loss rates have been found, it is added to a copy of
    the original GOESLightCurve object as goeslc.data.rad_loss_rate",
    where goeslc is the GOESLightCurve object.

    Parameters
    ----------
    goeslc : GOESLightCurve object

    Returns
    -------
    goeslc.data.rad_loss_rate : pandas.core.series.Series
                                Array of radiative loss rate of the
                                coronal soft X-ray-emitting plasma
                                across all wavelengths. [erg/s]

    Examples
    --------
    >>> from sunpy.lightcurve as lc
    >>> goeslc = lc.GOESLightCurve.create(time1, time2)
    >>> goeslc.data
                          xrsa   xrsb
    2014-01-01 00:00:00  7e-07  7e-06
    2014-01-01 00:00:02  7e-07  7e-06
    2014-01-01 00:00:04  7e-07  7e-06
    2014-01-01 00:00:06  7e-07  7e-06
    >>> goeslc_new = rad_loss_rate(goeslc)
    >>> goeslc_new.data
                          xrsa   xrsb  luminosity_xrsa  luminosity_xrsb
    2014-01-01 00:00:00  7e-07  7e-06     1.903523e+24     1.903523e+25
    2014-01-01 00:00:02  7e-07  7e-06     1.903523e+24     1.903523e+25
    2014-01-01 00:00:04  7e-07  7e-06     1.903523e+24     1.903523e+25
    2014-01-01 00:00:06  7e-07  7e-06     1.903523e+24     1.903523e+25

    """

    # Check that input argument is of correct type
    exceptions.check_goeslc(goeslc, varname="goeslc")

    # extract temperature and emission measure from GOESLightCurve
    # object and change type to that required by calc_rad_loss().
    # If GOESLightCurve object does not contain temperature and
    # emission measure, calculate using temp_em()
    try:
        temp = np.array(goeslc.data.temperature)
        em = np.array(goeslc.data.em)
    except AttributeError:
        goeslc_new = temp_em(goeslc)
        temp = np.array(goeslc_new.data.temperature)
        em = np.array(goeslc_new.data.em)
    else:
        goeslc_new = copy.deepcopy(goeslc)

    # Find radiative loss rate with calc_rad_loss()
    rad_loss_out = calc_rad_loss(temp, em)

    # Enter results into new version of GOES LightCurve Object
    goeslc_new.data["rad_loss_rate"] = rad_loss_out["rad_loss_rate"]

    return goeslc_new

def calc_rad_loss(temp, em, obstime=None):
    """
    Finds radiative loss rate of solar SXR plasma over all wavelengths.

    Extended Summary
    ----------------
    This function calculates the luminosity of solar coronal soft
    X-ray-emitting plasma across all wavelengths given an isothermal
    temperature and emission measure.  The units of the results are
    erg/s.  This function is based on calc_rad_loss.pro in SSW IDL.
    In addition, if obstime keyword is set, giving the times to which
    the temperature and emission measure values correspond, the
    radiated losses integrated over time are also calculated.

    Parameters
    ----------
    temp : numpy ndarray, dtype=float, units=[MK]
           Array containing the temperature of the coronal plasma at
           different times.
    em : numpy ndarray, dtype=float, units=[cm**-3]
         Array containing the emission measure of the coronal plasma
         at the same times corresponding to the temperatures in temp.
         Must be same length as temp
    obstime : numpy array, dtype=datetime64, optional
              array of measurement times to which temperature and
              emission measure values correspond.  Must be same length
              as temp and em.  If this keyword is set, the integrated
              radiated energy is calculated.

    Returns
    -------
    radlossrate : numpy ndarray, dtype=float
                  Array containing radiative loss rates of the coronal
                  plasma corresponding to temperatures and emission
                  measures in temp and em arrays.

    Notes
    -----
    This function calls a csv file containing a table of radiative loss
    rate per unit emission measure at various temperatures.  The
    appropriate values are then found via interpolation.
    This table was generated using CHIANTI atomic physics database
    employing the methods of Cox & Tucker (1969).  Coronal abundances,
    default density of 10**10 cm**-3, and ionization equilibrium of
    Mazzotta et al. (1998) were used.

    Examples
    --------
    >>> temp = np.array([11.28295376, 11.28295376])
    >>> em = np.array([4.78577516e+48, 4.78577516e+48])
    >>> rad_loss = calc_rad_loss(temp, em)
    >>> rad_loss
    array([  3.57994116e+26,   3.57994116e+26])
    """

    # Check inputs are correct
    exceptions.check_float(temp, varname="temp") # Check temp type
    exceptions.check_float(em, varname="em") # Check em type

    # Initialize lists to hold model data of temperature - rad loss rate
    # relationship read in from csv file
    modeltemp = [] # modelled temperature is in log_10 sapce in units of MK
    model_loss_rate = []

    # Read data from csv file into lists, being sure to skip commented
    # lines begining with "#"
    with open(os.path.join(INSTR_FILES_PATH, "chianti_rad_loss.csv"),
              "r") as csvfile:
        startline = dropwhile(lambda l: l.startswith("#"), csvfile)
        csvreader = csv.DictReader(startline, delimiter=";")
        for row in csvreader:
            modeltemp.append(float(row["temp_K"]))
            model_loss_rate.append(float(row["rad_loss_rate_per_em"]))
    modeltemp = np.array(modeltemp)
    model_loss_rate = np.array(model_loss_rate)
    # Ensure input values of flux ratio are within limits of model table
    if np.min(temp*1e6) < np.min(modeltemp) or
        np.max(temp*1e6) > np.max(modeltemp):
        raise ValueError("All values in temp must be within the range " +
                         "{0} - {1} MK.".format(np.min(modeltemp/1e6),
                                                np.max(modeltemp/1e6)))
    # Perform spline fit to model data to get temperatures for input
    # values of flux ratio
    spline = interpolate.splrep(modeltemp, model_loss_rate, s=0)
    rad_loss_rate = em * interpolate.splev(temp*1e6, spline, der=0)

    # If obstime keyword giving measurement times is set, calculate
    # radiative losses intergrated over time.
    if obstime is not None:
        dt = time_intervals(obstime)
        # Check that times are in chronological order
        if np.min(dt) <= 0:
            raise ValueError("times in obstime must be in " +
                             "chronological order.")
        rad_loss_int = np.sum(rad_loss_rate*dt)
        rad_loss_out = {"temperature":temp, "em":em,
                        "rad_loss_rate":rad_loss_rate, "time": obstime,
                        "rad_loss_int":rad_loss_int}
    else:
        rad_loss_out = {"temperature":temp, "em":em,
                        "rad_loss_rate":rad_loss_rate}

    return rad_loss_out

def xray_luminosity(goeslc):
    """
    Calculates and adds solar X-ray luminosity to a GOESLightCurve.

    Extended Summary
    ----------------
    This function calculates the solar X-ray luminosity in the
    GOES wavelength ranges (1-8 angstroms and 0.5-4 angstroms) based
    on the observed GOES fluxes.  The units of the results are erg/s.
    This is done by calling goes_lx().  This function assumes that the
    radiation is emitted isotropically, i.e. is distributed over a
    spherical surface area with a radius equal to the Sun-Earth
    distance.  Once the luminosity in each GOES passband is found,
    they are added to a copy of the original GOESLightCurve object as
    goeslc.data.luminosity_xrsa (for the 0.5-4 angstrom channel) and
    goeslc.data.luminosity_xrsb (for the 1-8 angstrom channel), where
    goeslc is the GOESLightCurve object.

    Parameters
    ----------
    goeslc : GOESLightCurve object

    Returns
    -------
    goeslc.data.luminosity_xrsa : pandas.core.series.Series
                                  Array of luminosity in the 0.5-4
                                  angstrom wavelength range [erg/s]
    goeslc.data.luminosity_xrsb : pandas.core.series.Series
                                  Array of luminosity in the 1-8
                                  angstrom wavelength range [erg/s]

    Examples
    --------
    >>> from sunpy.lightcurve as lc
    >>> goeslc = lc.GOESLightCurve.create(time1, time2)
    >>> goeslc.data
                          xrsa   xrsb
    2014-01-01 00:00:00  7e-07  7e-06
    2014-01-01 00:00:02  7e-07  7e-06
    2014-01-01 00:00:04  7e-07  7e-06
    2014-01-01 00:00:06  7e-07  7e-06
    >>> goeslc_new = xray_luminosity(goeslc)
    >>> goeslc_new.data
                          xrsa   xrsb  luminosity_xrsa  luminosity_xrsb
    2014-01-01 00:00:00  7e-07  7e-06     1.903523e+24     1.903523e+25
    2014-01-01 00:00:02  7e-07  7e-06     1.903523e+24     1.903523e+25
    2014-01-01 00:00:04  7e-07  7e-06     1.903523e+24     1.903523e+25
    2014-01-01 00:00:06  7e-07  7e-06     1.903523e+24     1.903523e+25
    """

    # Check that input argument is of correct type
    exceptions.check_goeslc(goeslc, varname="goeslc")

    # extract properties from GOESLightCurve object and change type to
    # that required by goes_chianti_em
    longflux = np.array(goeslc.data.xrsb)
    shortflux = np.array(goeslc.data.xrsa)
    date = str(goeslc.data.index[0])

    # Find temperature and emission measure with goes_chianti_tem
    lx_out = goes_lx(longflux, shortflux, date=date)

    # Enter results into new version of GOES LightCurve Object
    goeslc_new = copy.deepcopy(goeslc)
    goeslc_new.data["luminosity_xrsa"] = lx_out["shortlum"]
    goeslc_new.data["luminosity_xrsb"] = lx_out["longlum"]

    return goeslc_new

def goes_lx(longflux, shortflux, obstime=None, date=None):
    """Calculates solar X-ray luminosity in GOES wavelength ranges.

    Extended Summary
    ----------------
    This function calculates the X-ray luminosity from the Sun in the
    GOES wavelength ranges (1-8 angstroms and 0.5-4 angstroms) based
    on the observed GOES fluxes.  The units of the results are erg/s.
    The calculation is made by simply assuming that the radiation is
    emitted isotropically, i.e. is distributed over a spherical
    surface area with a radius equal to the Sun-Earth distance.

    Parameters
    ----------
    longflux : numpy ndarray, dtype=float
               Array containing the observed GOES/XRS long channel flux
    shortflux : numpy ndarray, dtype=float
                Array containing the observed GOES/XRS short channel
                flux
    obstime : numpy ndarray, dtype=datetime64, optional
              Measurement times corresponding to each long/short
              channel flux measurement.
    date : datetime object or valid date strng, optional
           Date at which measurements were taken.

    Returns
    -------
    longlum : numpy ndarray
              Array of luminosity in the long channel range
              (1-8 angstroms)
    shortlum : numpy ndarray
               Array of luminosity in the short channel range
               (0.5-4 angstroms)
    longlum_int : float
                  Long channel fluence, i.e. luminosity integrated
                  over time.
    shortlum_int : float
                   Short channel fluence, i.e. luminosity integrated
                   over time

    Notes
    -----
    This function calls goes_luminosity() to calculate luminosities.
    For more information on how this is done, see docstring of that
    function.

    Examples
    --------
    >>> longflux = np.array([7e-6,7e-6,7e-6,7e-6,7e-6,7e-6])
    >>> shortflux = np.array([7e-7,7e-7,7e-7,7e-7,7e-7,7e-7])
    >>> obstime = np.array(["2014-01-01 00:00:00",
                            "2014-01-01 00:00:02",
                            "2014-01-01 00:00:04"
                            "2014-01-01 00:00:06"
                            "2014-01-01 00:00:08"
                            "2014-01-01 00:00:10"]
                            dtype="datetime64[ms]")
    >>> lx_out = goes_lx(longflux, shortflux, obstime)
    >>> lx_out.longlum
    array([  1.98650769e+25,   1.98650769e+25,   1.98650769e+25,
             1.98650769e+25,   1.98650769e+25,   1.98650769e+25])
    >>> lx_out.shortlum
    array([  1.98650769e+24,   1.98650769e+24,   1.98650769e+24,
             1.98650769e+24,   1.98650769e+24,   1.98650769e+24])
    >>> lx_out.longlum_int
    2.0337865720138238e+26
    >>> lx_out.shortlum_int
    2.0337865720138235e+25

    """

    # Check inputs are of correct type
    exceptions.check_float(longflux, varname="longflux") # Check longflux type
    exceptions.check_float(shortflux, varname="shortflux") # Check shortflux type

    # Calculate X-ray luminosities
    longlum = calc_xraylum(longflux, date=date)
    shortlum = calc_xraylum(shortflux, date=date)

    # If obstime keyword giving measurement times is set, calculate
    # total energy radiated in the GOES bandpasses during the flare.
    if obstime is not None:
        dt = time_intervals(obstime)
        # Check that times are in chronological order
        if np.min(dt) <= 0:
            raise ValueError("times in obstime must be in "
                             "chronological order.")
        longlum_int = np.sum(longlum*dt)
        shortlum_int = np.sum(shortlum*dt)
        lx_out = {"longflux":longflux, "shortflux":shortflux,
                  "time":obstime, "longlum":longlum, "shortlum":shortlum,
                  "longlum_int":longlum_int, "shortlum_int":shortlum_int,
                  "dt":dt}
    else:
        lx_out = {"longflux":longflux, "shortflux":shortflux,
                  "longlum":longlum, "shortlum":shortlum,}

    return lx_out

def calc_xraylum(flux, date=None):
    """
    Calculates solar luminosity based on observed flux observed at 1AU.

    Extended Summary
    ----------------
    This function calculates the luminosity from the Sun based
    on observed flux in W/m**2.  The units of the results are erg/s.
    The calculation is made by simply assuming that the radiation is
    emitted isotropically, i.e. is distributed over a spherical
    surface area with a radius equal to the Sun-Earth distance.

    Parameters
    ----------
    flux : numpy ndarray
           Array containing the observed solar flux
    date : datetime object or valid date string, optional
           Used to calculate a more accurate Sun-Earth distance.

    Returns
    -------
    luminosity : numpy array
                Array of luminosity

    Notes
    -----
    To convert from W/m**2 to erg/s:
    1 W = 1 J/s = 10**7 erg/s
    1 W/m**2 = 4*pi * AU**2 * 10**7 erg/s, where AU is the Sun-Earth
    distance in metres.

    Examples
    --------
    >>> flux = np.array([7e-6,7e-6])
    >>> luminosity = goes_luminosity(flux, date="2014-04-21")
    >>> luminosity
    array([  1.98650769e+25,   1.98650769e+25])

    """
    exceptions.check_float(flux)
    if date is not None:
        date = exceptions.check_date(date)
        return 4 * np.pi * (sun.constants.au.value *
                            sun.sunearth_distance(t=date))**2 * 1e7 * flux
    else:
        return 4 * np.pi * (sun.constants.au.value)**2 * 1e7 * flux

def time_intervals(obstime):
    """
    Calculates time intervals between measurement times in seconds.

    Extended Summary
    ----------------
    This function calculates the time intervals between a series of
    measurement times for use in siple integration over time.
    Assume you have a series of times labelled t_1,...t_n.
    The start of the time bin for time t_i is defined as
    dt_i = (t_(i+1) - t_(i-1)) / 2
    i.e. from halfway between t_i and the previous time, t_(i-1), to
    halfway between t_i and the next time, t_(i+1).
    The time intervals for t_1 and t_n are special cases.  These are
    defined as
    dt_1 = (t_2 - t_1) / 2
    dt_(n-1) = (t_n - t_(n-1)) / 2

    Parameters
    ----------
    obstime : numpy ndarray, dtype=datetime64
              Array containing the time measurements.

    Returns
    -------
    dt : numpy array, dtype=float
         Array of time intervals in [s]

    Examples
    --------
    >>> obstime = np.array(["2014-01-01 00:00:00",
                            "2014-01-01 00:00:02",
                            "2014-01-01 00:00:04"
                            "2014-01-01 00:00:06"
                            "2014-01-01 00:00:08"
                            "2014-01-01 00:00:10"],
                            dtype="datetime64[ms]")
    >>> dt = time_intervals(obstime)
    >>> dt
    array([ 1.000,  2.000,  2.000,  2.000,  2.000,  1.000])

    """
    # check obstime is correct type and greater than min required length
    exceptions.check_datetime64(obstime, varname="obstime")
    if len(obstime) < 3:
        raise ValueError("obstime must have 3 or more elements.")
    obstime = obstime.astype("datetime64[ms]")  # convert to units of ms
    dt = (obstime[2:]-obstime[:-2]) / 2
    dt = np.insert(dt, 0, (obstime[1]-obstime[0])/2)
    dt = np.append(dt, (obstime[-1]-obstime[-2])/2)
    dt = dt.astype(float) / 1e3 # convert from [ms] to [s]
    return dt

def check_datetime64(test, varname="This variable"):
    """Raise Exception if test isn't numpy array of dtype datetime64.

    Parameters
    ----------
    test : variable to test
    varname : string, optional
              name of variable.  (Printed if exception is raised.)
    """
    if type(varname) is not str:
        varname = "This variable"
    if type(test) is not np.ndarray or test.dtype.type is not np.datetime64:
        raise TypeError(
            "{0} must be a numpy array of type datetime64.".format(varname))
<commit_msg>Removed goes_temp.py as it is still being developed should not have been included in this pull request,<commit_after><|endoftext|>


<commit_before>from __future__ import absolute_import

__authors__ = ["Russell Hewett, Stuart Mumford"]
__email__ = "stuart@mumford.me.uk"

import os
import glob
import urllib2

from collections import OrderedDict

import numpy as np

from sunpy.map import GenericMap
from sunpy.map.header import MapMeta
from sunpy.map.compositemap import CompositeMap
from sunpy.map.mapcube import MapCube

from sunpy.io.file_tools import read_file
from sunpy.io.header import FileHeader

from sunpy.util.net import download_file
from sunpy.util import expand_list
from sunpy.util.datatype_factory_base import RegisteredFactoryBase
from sunpy.util import Deprecated

__all__ = ['Map']

class Map(RegisteredFactoryBase):

    GenericWidgetType = GenericMap

    @classmethod
    def _read_file(cls, fname):
        # File gets read here.  This needs to be generic enough to seamlessly
        #call a fits file or a jpeg2k file, etc

        pairs  = read_file(fname)

        for i, pair in enumerate(pairs):
            filemeta = pair[1]
            filedata = pair[0]
            assert isinstance(filemeta, FileHeader)

            data = filedata
            meta = MapMeta(filemeta)
            pairs[i] = (data, meta)

        return pairs

    @classmethod
    def _parse_args(cls, *args, **kwargs):

        data_header_pairs = list()
        already_maps = list()

        # Account for nested lists of items
        args = expand_list(args)

        # For each of the arguments, handle each of the cases
        i = 0
        while i < len(args):

            arg = args[i]
            print arg

            # Data-header pair in a tuple
            if ((type(arg) in [tuple, list]) and
                 isinstance(arg[0],np.ndarray) and # or NDData or something else?
                 isinstance(arg[1],OrderedDict)): # FITSHeader, JP2kHeader, OrderedDict, dict?
                data_header_pairs.append(arg)

            # Data-header pair not in a tuple
            elif (isinstance(arg, np.ndarray) and # or NDData or something else?
                  isinstance(args[i+1],OrderedDict)): # FITSHeader, JP2kHeader, OrderedDict, dict?
                pair = (args[i], args[i+1])
                data_header_pairs.append(pair)
                i += 1 # an extra increment to account for the data-header pairing

            # File name
            elif (isinstance(arg,basestring) and
                  os.path.isfile(os.path.expanduser(arg))):
                path = os.path.expanduser(arg)
                pairs = cls._read_file(path)
                data_header_pairs += pairs

            # Directory
            elif (isinstance(arg,basestring) and
                  os.path.isdir(os.path.expanduser(arg))):
                path = os.path.expanduser(arg)
                files = [os.path.join(path, elem) for elem in os.listdir(path)]
                for afile in files:
                    data_header_pairs += cls._read_file(afile)

            # Glob
            elif (isinstance(arg,basestring) and '*' in arg):
                files = glob.glob( os.path.expanduser(arg) )
                for afile in files:
                    data_header_pairs += cls._read_file(afile)

            # Already a Map
            elif isinstance(arg, GenericMap):
                already_maps.append(arg)

            # A URL
            elif (isinstance(arg,basestring) and
                  urllib2.urlopen(arg)):
                default_dir = sunpy.config.get("downloads", "download_dir")
                url = arg
                path = download_file(url, default_dir)
                pairs = cls._read_file(path)
                data_header_pairs += pairs

            i += 1
        #TODO:
        # In the end, if there are aleady maps it should be put in the same
        # order as the input, currently they are not.
        return data_header_pairs, already_maps

    def __new__(cls, *args, **kwargs):

        # Hack to get around Python 2.x not backporting PEP 3102.
        composite = kwargs.pop('composite', False)
        cube = kwargs.pop('cube', False)

        if cls is Map:

            # Get list of data-header pairs, e.g., [(d1, h1), (d2, h2), ...]
            data_header_pairs, already_maps = cls._parse_args(*args, **kwargs)

            # Otherwise, each pair in the list gets built on its own
            new_maps = list()
            for pair in data_header_pairs:
                data, header = pair
                # Test to see which type of Map this pair is.  If none of the
                # registered Map types match, use a generic map.
                WidgetType = None
                for key in cls.registry:

                    if cls.registry[key](data, header, **kwargs):
                        WidgetType = key
                        break
                else:
                    WidgetType = cls.GenericWidgetType

                # Instantiate the new map.
                new_maps.append(WidgetType(data, header, **kwargs))

            new_maps += already_maps

            # If the list is meant to be a cube, instantiate a map cube
            if cube:
                return MapCube(new_maps, **kwargs)

            # If the list is meant to be a composite mape, instantiate one
            if composite:
                return CompositeMap(new_maps, **kwargs)

            # If there was only one map instantiated, return that, otherwise
            # return the list of them.
            return new_maps[0] if len(new_maps) == 1 else new_maps
        else:
            return super(Map, cls).__new__(cls, *args, **kwargs)

@Deprecated("Please use the new factory sunpy.Map")
def make_map(*args, **kwargs):
    __doc__ = Map.__doc__
    return Map(*args, **kwargs)

class InvalidMapInput(ValueError):
    """Exception to raise when input variable is not a Map instance and does
    not point to a valid Map input file."""
    pass

class InvalidMapType(ValueError):
    """Exception to raise when an invalid type of map is requested with make_map
    """
    pass

class NoMapsFound(ValueError):
    """Exception to raise when input does not point to any valid maps or files
    """
    pass<commit_msg>Added check that Map data is >1D<commit_after>from __future__ import absolute_import

__authors__ = ["Russell Hewett, Stuart Mumford"]
__email__ = "stuart@mumford.me.uk"

import os
import glob
import urllib2

from collections import OrderedDict

import numpy as np

from sunpy.map import GenericMap
from sunpy.map.header import MapMeta
from sunpy.map.compositemap import CompositeMap
from sunpy.map.mapcube import MapCube

from sunpy.io.file_tools import read_file
from sunpy.io.header import FileHeader

from sunpy.util.net import download_file
from sunpy.util import expand_list
from sunpy.util.datatype_factory_base import RegisteredFactoryBase
from sunpy.util import Deprecated

__all__ = ['Map']

class Map(RegisteredFactoryBase):

    GenericWidgetType = GenericMap

    @classmethod
    def _read_file(cls, fname):
        # File gets read here.  This needs to be generic enough to seamlessly
        #call a fits file or a jpeg2k file, etc

        pairs  = read_file(fname)
        new_pairs = []
        for i, pair in enumerate(pairs):
            filedata, filemeta = pair
            assert isinstance(filemeta, FileHeader)
            #This tests that the data is more than 1D
            if len(np.shape(filedata)) > 1:
                data = filedata
                meta = MapMeta(filemeta)
                new_pairs.append((data, meta))
        return new_pairs

    @classmethod
    def _parse_args(cls, *args, **kwargs):

        data_header_pairs = list()
        already_maps = list()

        # Account for nested lists of items
        args = expand_list(args)

        # For each of the arguments, handle each of the cases
        i = 0
        while i < len(args):

            arg = args[i]

            # Data-header pair in a tuple
            if ((type(arg) in [tuple, list]) and
                 isinstance(arg[0],np.ndarray) and # or NDData or something else?
                 isinstance(arg[1],OrderedDict)): # FITSHeader, JP2kHeader, OrderedDict, dict?
                data_header_pairs.append(arg)

            # Data-header pair not in a tuple
            elif (isinstance(arg, np.ndarray) and # or NDData or something else?
                  isinstance(args[i+1],OrderedDict)): # FITSHeader, JP2kHeader, OrderedDict, dict?
                pair = (args[i], args[i+1])
                data_header_pairs.append(pair)
                i += 1 # an extra increment to account for the data-header pairing

            # File name
            elif (isinstance(arg,basestring) and
                  os.path.isfile(os.path.expanduser(arg))):
                path = os.path.expanduser(arg)
                pairs = cls._read_file(path)
                data_header_pairs += pairs

            # Directory
            elif (isinstance(arg,basestring) and
                  os.path.isdir(os.path.expanduser(arg))):
                path = os.path.expanduser(arg)
                files = [os.path.join(path, elem) for elem in os.listdir(path)]
                for afile in files:
                    data_header_pairs += cls._read_file(afile)

            # Glob
            elif (isinstance(arg,basestring) and '*' in arg):
                files = glob.glob( os.path.expanduser(arg) )
                for afile in files:
                    data_header_pairs += cls._read_file(afile)

            # Already a Map
            elif isinstance(arg, GenericMap):
                already_maps.append(arg)

            # A URL
            elif (isinstance(arg,basestring) and
                  urllib2.urlopen(arg)):
                default_dir = sunpy.config.get("downloads", "download_dir")
                url = arg
                path = download_file(url, default_dir)
                pairs = cls._read_file(path)
                data_header_pairs += pairs

            i += 1
        #TODO:
        # In the end, if there are aleady maps it should be put in the same
        # order as the input, currently they are not.
        return data_header_pairs, already_maps

    def __new__(cls, *args, **kwargs):

        # Hack to get around Python 2.x not backporting PEP 3102.
        composite = kwargs.pop('composite', False)
        cube = kwargs.pop('cube', False)

        if cls is Map:

            # Get list of data-header pairs, e.g., [(d1, h1), (d2, h2), ...]
            data_header_pairs, already_maps = cls._parse_args(*args, **kwargs)

            # Otherwise, each pair in the list gets built on its own
            new_maps = list()
            for pair in data_header_pairs:
                data, header = pair
                # Test to see which type of Map this pair is.  If none of the
                # registered Map types match, use a generic map.
                WidgetType = None
                for key in cls.registry:

                    if cls.registry[key](data, header, **kwargs):
                        WidgetType = key
                        break
                else:
                    WidgetType = cls.GenericWidgetType

                # Instantiate the new map.
                new_maps.append(WidgetType(data, header, **kwargs))

            new_maps += already_maps

            # If the list is meant to be a cube, instantiate a map cube
            if cube:
                return MapCube(new_maps, **kwargs)

            # If the list is meant to be a composite mape, instantiate one
            if composite:
                return CompositeMap(new_maps, **kwargs)

            # If there was only one map instantiated, return that, otherwise
            # return the list of them.
            return new_maps[0] if len(new_maps) == 1 else new_maps
        else:
            return super(Map, cls).__new__(cls, *args, **kwargs)

@Deprecated("Please use the new factory sunpy.Map")
def make_map(*args, **kwargs):
    __doc__ = Map.__doc__
    return Map(*args, **kwargs)

class InvalidMapInput(ValueError):
    """Exception to raise when input variable is not a Map instance and does
    not point to a valid Map input file."""
    pass

class InvalidMapType(ValueError):
    """Exception to raise when an invalid type of map is requested with make_map
    """
    pass

class NoMapsFound(ValueError):
    """Exception to raise when input does not point to any valid maps or files
    """
    pass<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Sundry function and class decorators."""

from __future__ import print_function

import functools
import inspect
import textwrap
import types
import warnings

from sunpy.util.exceptions import SunpyDeprecationWarning
from sunpy.extern import six

__all__ = ['deprecated']

def deprecated(since, message='', name='', alternative=''):
    """
    Used to mark a function or class as deprecated.

    Parameters
    ------------
    since : str
        The release at which this API became deprecated.  This is
        required.

    message : str, optional
        Override the default deprecation message.  The format
        specifier ``func`` may be used for the name of the function,
        and ``alternative`` may be used in the deprecation message
        to insert the name of an alternative to the deprecated
        function. ``obj_type`` may be used to insert a friendly name
        for the type of object being deprecated.

    name : str, optional
        The name of the deprecated function or class; if not provided
        the name is automatically determined from the passed in
        function or class, though this is useful in the case of
        renamed functions, where the new function is just assigned to
        the name of the deprecated function.  For example::

            def new_function():
                ...
            oldFunction = new_function

    alternative : str, optional
        An alternative function or class name that the user may use in
        place of the deprecated object.  The deprecation warning will
        tell the user about this alternative if provided.

    """

    method_types = (classmethod, staticmethod, types.MethodType)

    def deprecate_doc(old_doc, message):
        """
        Returns a given docstring with a deprecation message prepended
        to it.
        """
        if not old_doc:
            old_doc = ''
        old_doc = textwrap.dedent(old_doc).strip('\n')
        new_doc = (('\n.. deprecated:: {since}'
                    '\n    {message}\n\n'.format(
                    **{'since': since, 'message': message.strip()})) + old_doc)
        if not old_doc:
            # This is to prevent a spurious 'unexpected unindent' warning from
            # docutils when the original docstring was blank.
            new_doc += r'\ '
        return new_doc

    def get_function(func):
        """
        Given a function or classmethod (or other function wrapper type), get
        the function object.
        """
        if isinstance(func, method_types):
            func = func.__func__
        return func

    def deprecate_function(func, message):
        """
        Returns a wrapped function that displays an
        ``AstropyDeprecationWarning`` when it is called.
        """

        if isinstance(func, method_types):
            func_wrapper = type(func)
        else:
            func_wrapper = lambda f: f

        func = get_function(func)

        def deprecated_func(*args, **kwargs):

            category = SunpyDeprecationWarning

            warnings.warn(message, category, stacklevel=2)

            return func(*args, **kwargs)

        # If this is an extension function, we can't call
        # functools.wraps on it, but we normally don't care.
        # This crazy way to get the type of a wrapper descriptor is
        # straight out of the Python 3.3 inspect module docs.
        if type(func) is not type(str.__dict__['__add__']):  # nopep8
            deprecated_func = functools.wraps(func)(deprecated_func)

        deprecated_func.__doc__ = deprecate_doc(
            deprecated_func.__doc__, message)

        return func_wrapper(deprecated_func)

    def deprecate_class(cls, message):
        """
        Returns a wrapper class with the docstrings updated and an
        __init__ function that will raise an
        ``AstropyDeprectationWarning`` warning when called.
        """
        # Creates a new class with the same name and bases as the
        # original class, but updates the dictionary with a new
        # docstring and a wrapped __init__ method.  __module__ needs
        # to be manually copied over, since otherwise it will be set
        # to *this* module (astropy.utils.misc).

        # This approach seems to make Sphinx happy (the new class
        # looks enough like the original class), and works with
        # extension classes (which functools.wraps does not, since
        # it tries to modify the original class).

        # We need to add a custom pickler or you'll get
        #     Can't pickle <class ..>: it's not found as ...
        # errors. Picklability is required for any class that is
        # documented by Sphinx.

        members = cls.__dict__.copy()

        members.update({
            '__doc__': deprecate_doc(cls.__doc__, message),
            '__init__': deprecate_function(get_function(cls.__init__),
                                           message),
        })

        return type(cls)(cls.__name__, cls.__bases__, members)

    def deprecate(obj, message=message, name=name, alternative=alternative):
        if isinstance(obj, type):
            obj_type_name = 'class'
        elif inspect.isfunction(obj):
            obj_type_name = 'function'
        elif inspect.ismethod(obj) or isinstance(obj, method_types):
            obj_type_name = 'method'
        else:
            obj_type_name = 'object'

        if not name:
            name = get_function(obj).__name__

        altmessage = ''
        if not message or type(message) is type(deprecate):
            message = ('The {func} {obj_type} is deprecated and may '
                       'be removed in a future version.')
            if alternative:
                altmessage = '\n        Use {} instead.'.format(alternative)

        message = ((message.format(**{
            'func': name,
            'name': name,
            'alternative': alternative,
            'obj_type': obj_type_name})) +
            altmessage)

        if isinstance(obj, type):
            return deprecate_class(obj, message)
        else:
            return deprecate_function(obj, message)

    if type(message) is type(deprecate):
        return deprecate(message)

    return deprecate

class add_common_docstring(object):
    """
    A function decorator that will append and/or prepend an addendum
    to the docstring of the target function.

    If a dictionary(d) is provided as parameter then format by .format().
    """
    def __init__(self, append='', prepend='', d={}):
        if any(d):
            append = append.format(**d)
            prepend = prepend.format(**d)
        self.append = append
        self.prepend = prepend

    def __call__(self, func):
        func.__doc__ = func.__doc__ if func.__doc__ else ''
        self.append = self.append if self.append else ''
        self.prepend = self.prepend if self.prepend else ''
        if self.append and isinstance(func.__doc__, six.string_types):
            func.__doc__ += self.append
        if self.prepend and isinstance(func.__doc__, six.string_types):
            func.__doc__ = self.prepend + func.__doc__
        return func
<commit_msg>made changes in init<commit_after># -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Sundry function and class decorators."""

from __future__ import print_function

import functools
import inspect
import textwrap
import types
import warnings

from sunpy.util.exceptions import SunpyDeprecationWarning
from sunpy.extern import six

__all__ = ['deprecated']

def deprecated(since, message='', name='', alternative=''):
    """
    Used to mark a function or class as deprecated.

    Parameters
    ------------
    since : str
        The release at which this API became deprecated.  This is
        required.

    message : str, optional
        Override the default deprecation message.  The format
        specifier ``func`` may be used for the name of the function,
        and ``alternative`` may be used in the deprecation message
        to insert the name of an alternative to the deprecated
        function. ``obj_type`` may be used to insert a friendly name
        for the type of object being deprecated.

    name : str, optional
        The name of the deprecated function or class; if not provided
        the name is automatically determined from the passed in
        function or class, though this is useful in the case of
        renamed functions, where the new function is just assigned to
        the name of the deprecated function.  For example::

            def new_function():
                ...
            oldFunction = new_function

    alternative : str, optional
        An alternative function or class name that the user may use in
        place of the deprecated object.  The deprecation warning will
        tell the user about this alternative if provided.

    """

    method_types = (classmethod, staticmethod, types.MethodType)

    def deprecate_doc(old_doc, message):
        """
        Returns a given docstring with a deprecation message prepended
        to it.
        """
        if not old_doc:
            old_doc = ''
        old_doc = textwrap.dedent(old_doc).strip('\n')
        new_doc = (('\n.. deprecated:: {since}'
                    '\n    {message}\n\n'.format(
                    **{'since': since, 'message': message.strip()})) + old_doc)
        if not old_doc:
            # This is to prevent a spurious 'unexpected unindent' warning from
            # docutils when the original docstring was blank.
            new_doc += r'\ '
        return new_doc

    def get_function(func):
        """
        Given a function or classmethod (or other function wrapper type), get
        the function object.
        """
        if isinstance(func, method_types):
            func = func.__func__
        return func

    def deprecate_function(func, message):
        """
        Returns a wrapped function that displays an
        ``AstropyDeprecationWarning`` when it is called.
        """

        if isinstance(func, method_types):
            func_wrapper = type(func)
        else:
            func_wrapper = lambda f: f

        func = get_function(func)

        def deprecated_func(*args, **kwargs):

            category = SunpyDeprecationWarning

            warnings.warn(message, category, stacklevel=2)

            return func(*args, **kwargs)

        # If this is an extension function, we can't call
        # functools.wraps on it, but we normally don't care.
        # This crazy way to get the type of a wrapper descriptor is
        # straight out of the Python 3.3 inspect module docs.
        if type(func) is not type(str.__dict__['__add__']):  # nopep8
            deprecated_func = functools.wraps(func)(deprecated_func)

        deprecated_func.__doc__ = deprecate_doc(
            deprecated_func.__doc__, message)

        return func_wrapper(deprecated_func)

    def deprecate_class(cls, message):
        """
        Returns a wrapper class with the docstrings updated and an
        __init__ function that will raise an
        ``AstropyDeprectationWarning`` warning when called.
        """
        # Creates a new class with the same name and bases as the
        # original class, but updates the dictionary with a new
        # docstring and a wrapped __init__ method.  __module__ needs
        # to be manually copied over, since otherwise it will be set
        # to *this* module (astropy.utils.misc).

        # This approach seems to make Sphinx happy (the new class
        # looks enough like the original class), and works with
        # extension classes (which functools.wraps does not, since
        # it tries to modify the original class).

        # We need to add a custom pickler or you'll get
        #     Can't pickle <class ..>: it's not found as ...
        # errors. Picklability is required for any class that is
        # documented by Sphinx.

        members = cls.__dict__.copy()

        members.update({
            '__doc__': deprecate_doc(cls.__doc__, message),
            '__init__': deprecate_function(get_function(cls.__init__),
                                           message),
        })

        return type(cls)(cls.__name__, cls.__bases__, members)

    def deprecate(obj, message=message, name=name, alternative=alternative):
        if isinstance(obj, type):
            obj_type_name = 'class'
        elif inspect.isfunction(obj):
            obj_type_name = 'function'
        elif inspect.ismethod(obj) or isinstance(obj, method_types):
            obj_type_name = 'method'
        else:
            obj_type_name = 'object'

        if not name:
            name = get_function(obj).__name__

        altmessage = ''
        if not message or type(message) is type(deprecate):
            message = ('The {func} {obj_type} is deprecated and may '
                       'be removed in a future version.')
            if alternative:
                altmessage = '\n        Use {} instead.'.format(alternative)

        message = ((message.format(**{
            'func': name,
            'name': name,
            'alternative': alternative,
            'obj_type': obj_type_name})) +
            altmessage)

        if isinstance(obj, type):
            return deprecate_class(obj, message)
        else:
            return deprecate_function(obj, message)

    if type(message) is type(deprecate):
        return deprecate(message)

    return deprecate

class add_common_docstring(object):
    """
    A function decorator that will append and/or prepend an addendum
    to the docstring of the target function.

    Parameters
    ----------

    append : `str`, optional
        A string to append to the end of the functions docstring.

    prepend : `str`, optional
        A string to prepend to the start of the functions docstring.
    """

    def __init__(self, append=None, prepend=None, **kwargs):
        if kwargs:
            append = append.format(**kwargs)
            prepend = prepend.format(**kwargs)
        self.append = append
        self.prepend = prepend

    def __call__(self, func):
        func.__doc__ = func.__doc__ if func.__doc__ else ''
        self.append = self.append if self.append else ''
        self.prepend = self.prepend if self.prepend else ''
        if self.append and isinstance(func.__doc__, six.string_types):
            func.__doc__ += self.append
        if self.prepend and isinstance(func.__doc__, six.string_types):
            func.__doc__ = self.prepend + func.__doc__
        return func
<|endoftext|>


<commit_before>#!/usr/bin/env python
from collections import deque
import copy
import numpy

# it may be simpler to represent these as an adjacency matrix directly, but i plan
# on using some big ass graphs, so we save some memory by doing it this way
class DirectedGraph(object):
    class Node(object):
        def __init__(self, value):
            self.outgoing = set()
            self.incoming = set()
            self.value = value

    def clear():
        self.nodes = {}

    def copy(self):
        return copy.deepcopy(self)

    def __init__(self):
        self.nodes = {}
        self.edges = {}
        self.edge_weights = {}

    @staticmethod
    def from_adjacency(nodes, adjM):
      g = DirectedGraph()

      nodesP = {}
      for k in nodes:
        nodesP[nodes[k]] = k
        g.add_node(k)

      for i in range(len(adjM)):
        for j in range(len(adjM)):
          if adjM[i,j] > 0:
            g.connect(nodesP[i], nodesP[j])
            g.set_weight(adjM[i,j], nodesP[i], nodesP[j])

      return g

    def add_node(self, node):
        self.nodes.setdefault(node, DirectedGraph.Node(node))

    def connect(self, src, dst, edgeValue=None, weight=1):
        if src not in self.nodes:
            self.nodes[src] = DirectedGraph.Node(src)
        if dst not in self.nodes:
            self.nodes[dst] = DirectedGraph.Node(dst)

        self.nodes[src].outgoing.add(dst)
        self.nodes[dst].incoming.add(src)

        if edgeValue != None:
            self.edges.setdefault((src,dst), set()).add(edgeValue)

    def set_weight(self, w, src, dst, edgeValue=None):
      self.connect(src,dst,edgeValue)
      self.edge_weights[(src,dst,edgeValue)] = w

    def get_weight(self, src, dst, edgeValue=None):
      if not self.connectedQ(src,dst):
        return 0

      if (src,dst,edgeValue) not in self.edge_weights:
        return 1

      return self.edge_weights[(src,dst,edgeValue)]

    def disconnect(self, src, dst):
      self.nodes[src].outgoing.remove(dst)
      self.nodes[dst].incoming.remove(src)

    def strip_edges_to(self, dst):
        n = self.nodes[dst]
        for i in n.incoming:
            ni = self.nodes[i]
            ni.outgoing = filter(lambda x: x != dst, ni.outgoing)
        n.incoming = set()

    def strip_edges_from(self, src):
        n = self.nodes[src]
        for i in n.outgoing:
            ni = self.nodes[i]
            ni.incoming = filter(lambda x: x != src, ni.incoming)
        n.outgoing = set()

    def remove_node(self, n):
        self.strip_edges_to(n)
        self.strip_edges_from(n)
        del self.nodes[n]

    def __contains__(self, n):
        return n in self.nodes

    def connectedQ(self, src, dst):
        src = self.nodes[src]
        for i in src.outgoing:
            if i == dst:
                return True
        return False

    def walk(self, src, direction='outgoing', depthfirst=False):
        src = self.nodes.setdefault(src, DirectedGraph.Node(src))
        q = deque([(src,0)])
        seen = set([src.value])

        while len(q) > 0:
            node,level = q.popleft()
            yield (node.value,level)
            seen.add(node.value)

            linked = set()
            if direction in ['either', 'outgoing']:
                linked = linked.union(node.outgoing)
            if direction in ['either', 'incoming']:
                linked = linked.union(node.incoming)

            for i in linked:
                if i in seen:
                    continue
                else:
                    seen.add(i)
                    if depthfirst:
                      q.appendleft((self.nodes[i], level+1))
                    else:
                      q.append((self.nodes[i],level+1))

    def within_distance(self, src, distance, direction='outgoing'):
        for n,l in self.walk(src, direction=direction):
            if l > distance:
                break
            yield n

    def stackwalk(self, src, direction='outgoing'):
        src = self.nodes.setdefault(src, DirectedGraph.Node(src))
        q = deque([[src.value]])
        seen = set([src.value])

        while len(q) > 0:
            stack = q.popleft()
            yield stack
            seen.add(stack[-1])
            n = self.nodes[stack[-1]]
            for i in (n.outgoing if direction == 'outgoing' else n.incoming):
                if i in seen:
                    continue
                else:
                    newstack = copy.copy(stack)
                    newstack.append(i)
                    seen.add(i)
                    q.append(newstack)

    def adjacency_matrix(self):
        ids = {}
        nid = 0
        for i in self.nodes.values():
            ids[i.value] = nid
            nid += 1

        m = numpy.zeros((nid,nid), dtype=float)
        for i in self.nodes.values():
            for j in i.outgoing:
                m[ids[i.value],ids[j]] = self.get_weight(i.value, j)

        return (ids, numpy.array(m,dtype=float))

    def _edge_count(self):
        ec = 0
        for n in self.nodes:
            ec += len(self.nodes[n].outgoing)
        return ec

    @property
    def exit_nodes(self):
      def _exit_nodes():
        for i in self.nodes.values():
          if len(i.outgoing) == 0:
            yield i

      return list(_exit_nodes())

    @property
    def start_node(self):
      for i in self.nodes.values():
        if len(i.incoming) == 0:
          return i

    def _write_dot_file(self,layout='dot'):
        import pydot
        import tempfile
        import os
        dotg = pydot.Dot('tmp', graph_type='digraph')

        dotnodes = {}
        for n in self.nodes:
            dotnodes[n] = pydot.Node(str(n))
            dotg.add_node(dotnodes[n])

        for n in self.nodes:
            for o in self.nodes[n].outgoing:
                if (n, o) in self.edges:
                    for e in self.edges[(n,o)]:
                        dotg.add_edge(pydot.Edge(dotnodes[n], dotnodes[o],label=e))
                else:
                    dotg.add_edge(pydot.Edge(dotnodes[n], dotnodes[o]))

        f = tempfile.NamedTemporaryFile(mode='w+b',delete=False)

        f.write(dotg.to_string())
        f.close()
        return f.name

    def visualize(self):
      '''
      visualize using xdot
      '''
      import os
      try:
        fname = self._write_dot_file()
        os.system('xdot --filter=%s %s' % (layout, fname))
      finally:
        os.unlink(fname)

    def _repr_svg_(self):
      '''
      visualize in IPython
      '''
      from IPython.display import SVG
      import os
      import tempfile
      fname = self._write_dot_file()
      f = tempfile.NamedTemporaryFile(mode='w+b', delete=False)
      f.close()
      os.system('dot -Tsvg -o %s %s' % (f.name, fname))
      svg = SVG(filename=f.name)
      os.unlink(f.name)
      os.unlink(fname)
      return svg.data

if __name__ == '__main__':
    from algorithms import *

    dg = DirectedGraph()
    dg.connect('a', 'b')
    dg.connect('b', 'c')
    dg.connect('b', 'd')
    dg.connect('d', 'a')

    print graph_string(dg)

    for node,level in dg.walk('a'):
        print "%s level[%d]" % (node,level)

    for stack in dg.stackwalk('a'):
        print stack

    cylic = list(find_cylic_nodes(dg, 'a'))
    print 'cyclic nodes: %s' % (cylic)

    print 'domination set: %s' % (dominate_sets(dg, 'a'),)
    print 'domination tree: %s' % (graph_string(domtree(dominate_sets(dg, 'a'))))
    print 'loop headers: %s' % (loop_headers(dg, dominate_sets(dg, 'a'), 'a'))
<commit_msg>added node coloring<commit_after>#!/usr/bin/env python
from collections import deque
import copy
import numpy

# it may be simpler to represent these as an adjacency matrix directly, but i plan
# on using some big ass graphs, so we save some memory by doing it this way
class DirectedGraph(object):
    class Node(object):
        def __init__(self, value):
            self.outgoing = set()
            self.incoming = set()
            self.value = value
            self.color = 'black'

    def clear():
        self.nodes = {}

    def get_color(self, node):
      return self.nodes[node].color

    def set_color(self, node, color):
      self.add_node(node)
      self.nodes[node].color = color

    def copy(self):
        return copy.deepcopy(self)

    def __init__(self):
        self.nodes = {}
        self.edges = {}
        self.edge_weights = {}

    @staticmethod
    def from_adjacency(nodes, adjM):
      g = DirectedGraph()

      nodesP = {}
      for k in nodes:
        nodesP[nodes[k]] = k
        g.add_node(k)

      for i in range(len(adjM)):
        for j in range(len(adjM)):
          if adjM[i,j] > 0:
            g.connect(nodesP[i], nodesP[j])
            g.set_weight(adjM[i,j], nodesP[i], nodesP[j])

      return g

    def add_node(self, node):
        self.nodes.setdefault(node, DirectedGraph.Node(node))

    def connect(self, src, dst, edgeValue=None, weight=1):
        if src not in self.nodes:
            self.nodes[src] = DirectedGraph.Node(src)
        if dst not in self.nodes:
            self.nodes[dst] = DirectedGraph.Node(dst)

        self.nodes[src].outgoing.add(dst)
        self.nodes[dst].incoming.add(src)

        if edgeValue != None:
            self.edges.setdefault((src,dst), set()).add(edgeValue)

    def set_weight(self, w, src, dst, edgeValue=None):
      self.connect(src,dst,edgeValue)
      self.edge_weights[(src,dst,edgeValue)] = w

    def get_weight(self, src, dst, edgeValue=None):
      if not self.connectedQ(src,dst):
        return 0

      if (src,dst,edgeValue) not in self.edge_weights:
        return 1

      return self.edge_weights[(src,dst,edgeValue)]

    def disconnect(self, src, dst):
      self.nodes[src].outgoing.remove(dst)
      self.nodes[dst].incoming.remove(src)

    def strip_edges_to(self, dst):
        n = self.nodes[dst]
        for i in n.incoming:
            ni = self.nodes[i]
            ni.outgoing = filter(lambda x: x != dst, ni.outgoing)
        n.incoming = set()

    def strip_edges_from(self, src):
        n = self.nodes[src]
        for i in n.outgoing:
            ni = self.nodes[i]
            ni.incoming = filter(lambda x: x != src, ni.incoming)
        n.outgoing = set()

    def remove_node(self, n):
        self.strip_edges_to(n)
        self.strip_edges_from(n)
        del self.nodes[n]

    def __contains__(self, n):
        return n in self.nodes

    def connectedQ(self, src, dst):
        src = self.nodes[src]
        for i in src.outgoing:
            if i == dst:
                return True
        return False

    def walk(self, src, direction='outgoing', depthfirst=False):
        src = self.nodes.setdefault(src, DirectedGraph.Node(src))
        q = deque([(src,0)])
        seen = set([src.value])

        while len(q) > 0:
            node,level = q.popleft()
            yield (node.value,level)
            seen.add(node.value)

            linked = set()
            if direction in ['either', 'outgoing']:
                linked = linked.union(node.outgoing)
            if direction in ['either', 'incoming']:
                linked = linked.union(node.incoming)

            for i in linked:
                if i in seen:
                    continue
                else:
                    seen.add(i)
                    if depthfirst:
                      q.appendleft((self.nodes[i], level+1))
                    else:
                      q.append((self.nodes[i],level+1))

    def within_distance(self, src, distance, direction='outgoing'):
        for n,l in self.walk(src, direction=direction):
            if l > distance:
                break
            yield n

    def stackwalk(self, src, direction='outgoing'):
        src = self.nodes.setdefault(src, DirectedGraph.Node(src))
        q = deque([[src.value]])
        seen = set([src.value])

        while len(q) > 0:
            stack = q.popleft()
            yield stack
            seen.add(stack[-1])
            n = self.nodes[stack[-1]]
            for i in (n.outgoing if direction == 'outgoing' else n.incoming):
                if i in seen:
                    continue
                else:
                    newstack = copy.copy(stack)
                    newstack.append(i)
                    seen.add(i)
                    q.append(newstack)

    def adjacency_matrix(self):
        ids = {}
        nid = 0
        for i in self.nodes.values():
            ids[i.value] = nid
            nid += 1

        m = numpy.zeros((nid,nid), dtype=float)
        for i in self.nodes.values():
            for j in i.outgoing:
                m[ids[i.value],ids[j]] = self.get_weight(i.value, j)

        return (ids, numpy.array(m,dtype=float))

    def _edge_count(self):
        ec = 0
        for n in self.nodes:
            ec += len(self.nodes[n].outgoing)
        return ec

    @property
    def exit_nodes(self):
      def _exit_nodes():
        for i in self.nodes.values():
          if len(i.outgoing) == 0:
            yield i

      return list(_exit_nodes())

    @property
    def start_node(self):
      for i in self.nodes.values():
        if len(i.incoming) == 0:
          return i

    def _write_dot_file(self,layout='dot'):
        import pydot
        import tempfile
        import os
        dotg = pydot.Dot('tmp', graph_type='digraph')

        dotnodes = {}
        for n in self.nodes:
            dotnodes[n] = pydot.Node(str(n), color=self.get_color(n))
            dotg.add_node(dotnodes[n])

        for n in self.nodes:
            for o in self.nodes[n].outgoing:
                if (n, o) in self.edges:
                    for e in self.edges[(n,o)]:
                        dotg.add_edge(pydot.Edge(dotnodes[n], dotnodes[o],label=e))
                else:
                    dotg.add_edge(pydot.Edge(dotnodes[n], dotnodes[o]))

        f = tempfile.NamedTemporaryFile(mode='w+b',delete=False)

        f.write(dotg.to_string())
        f.close()
        return f.name

    def visualize(self):
      '''
      visualize using xdot
      '''
      import os
      try:
        fname = self._write_dot_file()
        os.system('xdot --filter=%s %s' % (layout, fname))
      finally:
        os.unlink(fname)

    def _repr_svg_(self):
      '''
      visualize in IPython
      '''
      from IPython.display import SVG
      import os
      import tempfile
      fname = self._write_dot_file()
      f = tempfile.NamedTemporaryFile(mode='w+b', delete=False)
      f.close()
      os.system('dot -Tsvg -o %s %s' % (f.name, fname))
      svg = SVG(filename=f.name)
      os.unlink(f.name)
      os.unlink(fname)
      return svg.data

if __name__ == '__main__':
    from algorithms import *

    dg = DirectedGraph()
    dg.connect('a', 'b')
    dg.connect('b', 'c')
    dg.connect('b', 'd')
    dg.connect('d', 'a')

    print graph_string(dg)

    for node,level in dg.walk('a'):
        print "%s level[%d]" % (node,level)

    for stack in dg.stackwalk('a'):
        print stack

    cylic = list(find_cylic_nodes(dg, 'a'))
    print 'cyclic nodes: %s' % (cylic)

    print 'domination set: %s' % (dominate_sets(dg, 'a'),)
    print 'domination tree: %s' % (graph_string(domtree(dominate_sets(dg, 'a'))))
    print 'loop headers: %s' % (loop_headers(dg, dominate_sets(dg, 'a'), 'a'))
<|endoftext|>


<commit_before>from collections import OrderedDict

from sympy import Add, Mul, collect, collect_const

from devito.ir import DummyEq, Cluster, Scope
from devito.symbolics import (count, estimate_cost, q_xop, q_leaf, retrieve_scalars,
                              retrieve_terminals, yreplace)
from devito.tools import DAG, ReducerMap, split
from devito.types import Dimension, Symbol

__all__ = ['collect_nested', 'common_subexprs_elimination', 'make_is_time_invariant']

def collect_nested(expr):
    """
    Collect numeric coefficients, trascendental functions, and symbolic powers,
    across all levels of the expression tree.

    The collection gives precedence to (in order of importance):

        1) Trascendental functions,
        2) Symbolic powers,
        3) Numeric coefficients.

    Parameters
    ----------
    expr : expr-like
        The expression to be factorized.
    """

    def run(expr):
        # Return semantic (rebuilt expression, factorization candidates)

        if expr.is_Number:
            return expr, {'coeffs': expr}
        elif expr.is_Function:
            return expr, {'funcs': expr}
        elif expr.is_Pow:
            return expr, {'pows': expr}
        elif expr.is_Symbol or expr.is_Indexed or expr.is_Atom:
            return expr, {}
        elif expr.is_Add:
            args, candidates = zip(*[run(arg) for arg in expr.args])
            candidates = ReducerMap.fromdicts(*candidates)

            funcs = candidates.getall('funcs', [])
            pows = candidates.getall('pows', [])
            coeffs = candidates.getall('coeffs', [])

            # Functions/Pows are collected first, coefficients afterwards
            # Note: below we use sets, but SymPy will ensure determinism
            args = set(args)
            w_funcs = {i for i in args if any(j in funcs for j in i.args)}
            args -= w_funcs
            w_pows = {i for i in args if any(j in pows for j in i.args)}
            args -= w_pows
            w_coeffs = {i for i in args if any(j in coeffs for j in i.args)}
            args -= w_coeffs

            # Collect common funcs
            w_funcs = collect(expr.func(*w_funcs), funcs, evaluate=False)
            try:
                w_funcs = Add(*[Mul(k, collect_const(v)) for k, v in w_funcs.items()])
            except AttributeError:
                assert w_funcs == 0

            # Collect common pows
            w_pows = collect(expr.func(*w_pows), pows, evaluate=False)
            try:
                w_pows = Add(*[Mul(k, collect_const(v)) for k, v in w_pows.items()])
            except AttributeError:
                assert w_pows == 0

            # Collect common temporaries (r0, r1, ...)
            w_coeffs = collect(expr.func(*w_coeffs), tuple(retrieve_scalars(expr)),
                               evaluate=False)
            try:
                w_coeffs = Add(*[Mul(k, collect_const(v)) for k, v in w_coeffs.items()])
            except AttributeError:
                assert w_coeffs == 0

            # Collect common coefficients
            w_coeffs = collect_const(w_coeffs)

            rebuilt = Add(w_funcs, w_pows, w_coeffs, *args)

            return rebuilt, {}
        elif expr.is_Mul:
            args, candidates = zip(*[run(arg) for arg in expr.args])

            # Always collect coefficients
            rebuilt = collect_const(expr.func(*args))
            try:
                if rebuilt.args:
                    # Note: Mul(*()) -> 1, and since sympy.S.Zero.args == (),
                    # the `if` prevents turning 0 into 1
                    rebuilt = Mul(*rebuilt.args)
            except AttributeError:
                pass

            return rebuilt, ReducerMap.fromdicts(*candidates)
        elif expr.is_Equality:
            args, candidates = zip(*[run(expr.lhs), run(expr.rhs)])
            return expr.func(*args, evaluate=False), ReducerMap.fromdicts(*candidates)
        else:
            args, candidates = zip(*[run(arg) for arg in expr.args])
            return expr.func(*args), ReducerMap.fromdicts(*candidates)

    return run(expr)[0]

def make_is_time_invariant(context):
    """
    Given an ordered list of expressions, returns a callable that finds out whether
    a given expression is time invariant or not.
    """
    mapper = OrderedDict([(i.lhs, i) for i in _makeit_ssa(context)])

    def is_time_invariant(mapper, expr):
        if any(isinstance(i, Dimension) and i.is_Time for i in expr.free_symbols):
            return False

        queue = [expr.rhs if expr.is_Equality else expr]
        seen = set()
        while queue:
            item = queue.pop()
            nodes = set()
            for i in retrieve_terminals(item):
                if i in seen:
                    # Already inspected, nothing more can be inferred
                    continue
                elif any(isinstance(j, Dimension) and j.is_Time for j in i.free_symbols):
                    # Definitely not time-invariant
                    return False
                elif i in mapper:
                    # Go on with the search
                    nodes.add(i)
                elif isinstance(i, Dimension):
                    # Go on with the search, as `i` is not a time dimension
                    pass
                elif not i.function.is_DiscreteFunction:
                    # It didn't come from the outside and it's not in `mapper`, so
                    # cannot determine if time-invariant; assume time-varying then
                    return False
                seen.add(i)
            queue.extend([mapper[i].rhs for i in nodes])
        return True

    callback = lambda i: is_time_invariant(mapper, i)

    return callback

def common_subexprs_elimination(maybe_exprs, make, mode='default'):
    """
    Perform common sub-expressions elimination, or CSE.

    Note: the output is guaranteed to be topologically sorted.

    Parameters
    ----------
    maybe_exprs : expr-like or list of expr-like or Cluster
        One or more expressions to which CSE is applied.
    make : callable
        Build symbols to store temporary, redundant values.
    mode : str, optional
        The CSE algorithm applied. Accepted: ['default'].
    """

    # Note: not defaulting to SymPy's CSE() function for three reasons:
    # - it also captures array index access functions (eg, i+1 in A[i+1] and B[i+1]);
    # - it sometimes "captures too much", losing factorization opportunities;
    # - very slow
    # TODO: a second "sympy" mode will be provided, relying on SymPy's CSE() but
    # also ensuring some sort of post-processing
    assert mode == 'default'  # Only supported mode ATM

    # Just for flexibility, accept either Clusters or exprs
    if isinstance(maybe_exprs, Cluster):
        cluster = maybe_exprs
        processed = list(cluster.exprs)
        scope = cluster.scope
    else:
        processed = list(maybe_exprs)
        scope = Scope(maybe_exprs)

    # Some sub-expressions aren't really "common" -- that's the case of Dimension-
    # independent data dependences. For example:
    #
    # a[i] = ...
    # ... = ... a[i] + 1 ...
    # ... = ...
    # ... = ... a[i] + 1 ...
    #
    # `a[i] + 1` will be excluded, as there's a Dimension-independent data dependence
    # involving `a`
    exclude = {i.source.indexed for i in scope.d_all.independent()}

    mapped = []
    while True:
        # Detect redundancies
        counted = count(mapped + processed, q_xop).items()
        targets = OrderedDict([(k, estimate_cost(k, True)) for k, v in counted if v > 1])

        # Rule out Dimension-independent data dependencies
        targets = OrderedDict([(k, v) for k, v in targets.items()
                               if not k.free_symbols & exclude])

        if not targets:
            break

        # Create temporaries
        hit = max(targets.values())
        picked = [k for k, v in targets.items() if v == hit]
        mapper = OrderedDict([(e, make()) for i, e in enumerate(picked)])

        # Apply replacements
        processed = [e.xreplace(mapper) for e in processed]
        mapped = [e.xreplace(mapper) for e in mapped]
        mapped = [DummyEq(v, k) for k, v in reversed(list(mapper.items()))] + mapped

        # Prepare for the next round
        for k in picked:
            targets.pop(k)
    processed = mapped + processed

    # At this point we may have useless temporaries (e.g., r0=r1). Let's drop them
    processed = _compact_temporaries(processed)

    # Perform topological sorting so that reads-after-writes are honored
    processed = _topological_sort(processed)

    return processed

# Private functions

def _makeit_ssa(exprs):
    """
    Convert an iterable of Eqs into Static Single Assignment (SSA) form.
    """
    # Identify recurring LHSs
    seen = {}
    for i, e in enumerate(exprs):
        seen.setdefault(e.lhs, []).append(i)
    # Optimization: don't waste time reconstructing stuff if already in SSA form
    if all(len(i) == 1 for i in seen.values()):
        return exprs
    # SSA conversion
    c = 0
    mapper = {}
    processed = []
    for i, e in enumerate(exprs):
        where = seen[e.lhs]
        rhs = e.rhs.xreplace(mapper)
        if len(where) > 1:
            needssa = e.is_Scalar or where[-1] != i
            lhs = Symbol(name='ssa%d' % c, dtype=e.dtype) if needssa else e.lhs
            if e.is_Increment:
                # Turn AugmentedAssignment into Assignment
                processed.append(e.func(lhs, mapper[e.lhs] + rhs, is_Increment=False))
            else:
                processed.append(e.func(lhs, rhs))
            mapper[e.lhs] = lhs
            c += 1
        else:
            processed.append(e.func(e.lhs, rhs))
    return processed

def _compact_temporaries(exprs):
    """
    Drop temporaries consisting of isolated symbols.
    """
    # First of all, convert to SSA
    exprs = _makeit_ssa(exprs)

    # What's gonna be dropped
    mapper = {e.lhs: e.rhs for e in exprs
              if e.lhs.is_Symbol and (q_leaf(e.rhs) or e.rhs.is_Function)}

    processed = []
    for e in exprs:
        if e.lhs not in mapper:
            # The temporary is retained, and substitutions may be applied
            handle, _ = yreplace(e, mapper, repeat=True)
            assert len(handle) == 1
            processed.extend(handle)

    return processed

def _topological_sort(exprs):
    """
    Topologically sort the temporaries in a list of equations.
    """
    mapper = {e.lhs: e for e in exprs}
    assert len(mapper) == len(exprs)  # Expect SSA

    # Build DAG and topologically-sort temporaries
    temporaries, tensors = split(exprs, lambda e: not e.lhs.is_Indexed)
    dag = DAG(nodes=temporaries)
    for e in temporaries:
        for r in retrieve_terminals(e.rhs):
            if r not in mapper:
                continue
            elif mapper[r] is e:
                # Avoid cyclic dependences, such as
                # Eq(f, f + 1)
                continue
            elif r.is_Indexed:
                # Only scalars enforce an ordering
                continue
            else:
                dag.add_edge(mapper[r], e, force_add=True)
    processed = dag.topological_sort()

    # Append tensor equations at the end in user-provided order
    processed.extend(tensors)

    return processed
<commit_msg>dse: Drop now useless _toposort()<commit_after>from collections import OrderedDict

from sympy import Add, Mul, collect, collect_const

from devito.ir import DummyEq, Cluster, Scope
from devito.symbolics import (count, estimate_cost, q_xop, q_leaf, retrieve_scalars,
                              retrieve_terminals, yreplace)
from devito.tools import ReducerMap
from devito.types import Dimension, Symbol

__all__ = ['collect_nested', 'common_subexprs_elimination', 'make_is_time_invariant']

def collect_nested(expr):
    """
    Collect numeric coefficients, trascendental functions, and symbolic powers,
    across all levels of the expression tree.

    The collection gives precedence to (in order of importance):

        1) Trascendental functions,
        2) Symbolic powers,
        3) Numeric coefficients.

    Parameters
    ----------
    expr : expr-like
        The expression to be factorized.
    """

    def run(expr):
        # Return semantic (rebuilt expression, factorization candidates)

        if expr.is_Number:
            return expr, {'coeffs': expr}
        elif expr.is_Function:
            return expr, {'funcs': expr}
        elif expr.is_Pow:
            return expr, {'pows': expr}
        elif expr.is_Symbol or expr.is_Indexed or expr.is_Atom:
            return expr, {}
        elif expr.is_Add:
            args, candidates = zip(*[run(arg) for arg in expr.args])
            candidates = ReducerMap.fromdicts(*candidates)

            funcs = candidates.getall('funcs', [])
            pows = candidates.getall('pows', [])
            coeffs = candidates.getall('coeffs', [])

            # Functions/Pows are collected first, coefficients afterwards
            # Note: below we use sets, but SymPy will ensure determinism
            args = set(args)
            w_funcs = {i for i in args if any(j in funcs for j in i.args)}
            args -= w_funcs
            w_pows = {i for i in args if any(j in pows for j in i.args)}
            args -= w_pows
            w_coeffs = {i for i in args if any(j in coeffs for j in i.args)}
            args -= w_coeffs

            # Collect common funcs
            w_funcs = collect(expr.func(*w_funcs), funcs, evaluate=False)
            try:
                w_funcs = Add(*[Mul(k, collect_const(v)) for k, v in w_funcs.items()])
            except AttributeError:
                assert w_funcs == 0

            # Collect common pows
            w_pows = collect(expr.func(*w_pows), pows, evaluate=False)
            try:
                w_pows = Add(*[Mul(k, collect_const(v)) for k, v in w_pows.items()])
            except AttributeError:
                assert w_pows == 0

            # Collect common temporaries (r0, r1, ...)
            w_coeffs = collect(expr.func(*w_coeffs), tuple(retrieve_scalars(expr)),
                               evaluate=False)
            try:
                w_coeffs = Add(*[Mul(k, collect_const(v)) for k, v in w_coeffs.items()])
            except AttributeError:
                assert w_coeffs == 0

            # Collect common coefficients
            w_coeffs = collect_const(w_coeffs)

            rebuilt = Add(w_funcs, w_pows, w_coeffs, *args)

            return rebuilt, {}
        elif expr.is_Mul:
            args, candidates = zip(*[run(arg) for arg in expr.args])

            # Always collect coefficients
            rebuilt = collect_const(expr.func(*args))
            try:
                if rebuilt.args:
                    # Note: Mul(*()) -> 1, and since sympy.S.Zero.args == (),
                    # the `if` prevents turning 0 into 1
                    rebuilt = Mul(*rebuilt.args)
            except AttributeError:
                pass

            return rebuilt, ReducerMap.fromdicts(*candidates)
        elif expr.is_Equality:
            args, candidates = zip(*[run(expr.lhs), run(expr.rhs)])
            return expr.func(*args, evaluate=False), ReducerMap.fromdicts(*candidates)
        else:
            args, candidates = zip(*[run(arg) for arg in expr.args])
            return expr.func(*args), ReducerMap.fromdicts(*candidates)

    return run(expr)[0]

def make_is_time_invariant(context):
    """
    Given an ordered list of expressions, returns a callable that finds out whether
    a given expression is time invariant or not.
    """
    mapper = OrderedDict([(i.lhs, i) for i in _makeit_ssa(context)])

    def is_time_invariant(mapper, expr):
        if any(isinstance(i, Dimension) and i.is_Time for i in expr.free_symbols):
            return False

        queue = [expr.rhs if expr.is_Equality else expr]
        seen = set()
        while queue:
            item = queue.pop()
            nodes = set()
            for i in retrieve_terminals(item):
                if i in seen:
                    # Already inspected, nothing more can be inferred
                    continue
                elif any(isinstance(j, Dimension) and j.is_Time for j in i.free_symbols):
                    # Definitely not time-invariant
                    return False
                elif i in mapper:
                    # Go on with the search
                    nodes.add(i)
                elif isinstance(i, Dimension):
                    # Go on with the search, as `i` is not a time dimension
                    pass
                elif not i.function.is_DiscreteFunction:
                    # It didn't come from the outside and it's not in `mapper`, so
                    # cannot determine if time-invariant; assume time-varying then
                    return False
                seen.add(i)
            queue.extend([mapper[i].rhs for i in nodes])
        return True

    callback = lambda i: is_time_invariant(mapper, i)

    return callback

def common_subexprs_elimination(maybe_exprs, make, mode='default'):
    """
    Perform common sub-expressions elimination, or CSE.

    Note: the output is guaranteed to be topologically sorted.

    Parameters
    ----------
    maybe_exprs : expr-like or list of expr-like or Cluster
        One or more expressions to which CSE is applied.
    make : callable
        Build symbols to store temporary, redundant values.
    mode : str, optional
        The CSE algorithm applied. Accepted: ['default'].
    """

    # Note: not defaulting to SymPy's CSE() function for three reasons:
    # - it also captures array index access functions (eg, i+1 in A[i+1] and B[i+1]);
    # - it sometimes "captures too much", losing factorization opportunities;
    # - very slow
    # TODO: a second "sympy" mode will be provided, relying on SymPy's CSE() but
    # also ensuring some sort of post-processing
    assert mode == 'default'  # Only supported mode ATM

    # Just for flexibility, accept either Clusters or exprs
    if isinstance(maybe_exprs, Cluster):
        cluster = maybe_exprs
        processed = list(cluster.exprs)
        scope = cluster.scope
    else:
        processed = list(maybe_exprs)
        scope = Scope(maybe_exprs)

    # Some sub-expressions aren't really "common" -- that's the case of Dimension-
    # independent data dependences. For example:
    #
    # a[i] = ...
    # ... = ... a[i] + 1 ...
    # ... = ...
    # ... = ... a[i] + 1 ...
    #
    # `a[i] + 1` will be excluded, as there's a Dimension-independent data dependence
    # involving `a`
    exclude = {i.source.indexed for i in scope.d_all.independent()}

    mapped = []
    while True:
        # Detect redundancies
        counted = count(mapped + processed, q_xop).items()
        targets = OrderedDict([(k, estimate_cost(k, True)) for k, v in counted if v > 1])

        # Rule out Dimension-independent data dependencies
        targets = OrderedDict([(k, v) for k, v in targets.items()
                               if not k.free_symbols & exclude])

        if not targets:
            break

        # Create temporaries
        hit = max(targets.values())
        picked = [k for k, v in targets.items() if v == hit]
        mapper = OrderedDict([(e, make()) for i, e in enumerate(picked)])

        # Apply replacements
        processed = [e.xreplace(mapper) for e in processed]
        mapped = [e.xreplace(mapper) for e in mapped]
        mapped = [DummyEq(v, k) for k, v in reversed(list(mapper.items()))] + mapped

        # Prepare for the next round
        for k in picked:
            targets.pop(k)
    processed = mapped + processed

    # At this point we may have useless temporaries (e.g., r0=r1). Let's drop them
    processed = _compact_temporaries(processed)

    return processed

# Private functions

def _makeit_ssa(exprs):
    """
    Convert an iterable of Eqs into Static Single Assignment (SSA) form.
    """
    # Identify recurring LHSs
    seen = {}
    for i, e in enumerate(exprs):
        seen.setdefault(e.lhs, []).append(i)
    # Optimization: don't waste time reconstructing stuff if already in SSA form
    if all(len(i) == 1 for i in seen.values()):
        return exprs
    # SSA conversion
    c = 0
    mapper = {}
    processed = []
    for i, e in enumerate(exprs):
        where = seen[e.lhs]
        rhs = e.rhs.xreplace(mapper)
        if len(where) > 1:
            needssa = e.is_Scalar or where[-1] != i
            lhs = Symbol(name='ssa%d' % c, dtype=e.dtype) if needssa else e.lhs
            if e.is_Increment:
                # Turn AugmentedAssignment into Assignment
                processed.append(e.func(lhs, mapper[e.lhs] + rhs, is_Increment=False))
            else:
                processed.append(e.func(lhs, rhs))
            mapper[e.lhs] = lhs
            c += 1
        else:
            processed.append(e.func(e.lhs, rhs))
    return processed

def _compact_temporaries(exprs):
    """
    Drop temporaries consisting of isolated symbols.
    """
    # First of all, convert to SSA
    exprs = _makeit_ssa(exprs)

    # What's gonna be dropped
    mapper = {e.lhs: e.rhs for e in exprs
              if e.lhs.is_Symbol and (q_leaf(e.rhs) or e.rhs.is_Function)}

    processed = []
    for e in exprs:
        if e.lhs not in mapper:
            # The temporary is retained, and substitutions may be applied
            handle, _ = yreplace(e, mapper, repeat=True)
            assert len(handle) == 1
            processed.extend(handle)

    return processed
<|endoftext|>


<commit_before>"""User API to specify relationals."""

import sympy

__all__ = ['Le', 'Lt', 'Ge', 'Gt']

class Le(sympy.Le):
    """
    A less-than ("<=") relation between two objects, the left-hand side and the
    right-hand side. Can be used to build conditionals but not directly to
    construct an Operator.

    The left-hand side may be a Function or a SparseFunction. The right-hand
    side may be any arbitrary expressions with numbers, Dimensions, Constants,
    Functions and SparseFunctions as operands.

    Parameters
    ----------
    lhs : Function or SparseFunction
        The left-hand side.
    rhs : expr-like, optional
        The right-hand side. Defaults to 0.
    subdomain : SubDomain, optional
        To restrict the evalaution of the relation to a particular sub-region in the
        computational domain.

    Examples
    --------
    Le may be used to express a relation (e.g. in a Subdomain).

    >>> from devito import Grid, Function, ConditionalDimension, Eq, Operator
    >>> from devito.types import Le
    >>> grid = Grid(shape=(8, 8))
    >>> g = Function(name='g', shape=grid.shape, dimensions=grid.dimensions)
    >>> x, y = grid.dimensions
    >>> cond = Le(g, 1)
    >>> ci = ConditionalDimension(name='ci', parent=y, condition=cond)
    >>> f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
    >>> op = Operator(Eq(g, f))
    """

    def __init__(self, lhs, rhs=0, subdomain=None, **kwargs):
        obj = sympy.Le.__new__(self, lhs, rhs, **kwargs)
        obj._subdomain = subdomain
        return obj

    @property
    def subdomain(self):
        """The SubDomain in which the Le is defined."""
        return self._subdomain

class Lt(sympy.Lt):
    """
    A strict-less-than ("<") relation between two objects, the left-hand side and the
    right-hand side. Can be used to build conditionals but not directly to
    construct an Operator.

    The left-hand side may be a Function or a SparseFunction. The right-hand
    side may be any arbitrary expressions with numbers, Dimensions, Constants,
    Functions and SparseFunctions as operands.

    Parameters
    ----------
    lhs : Function or SparseFunction
        The left-hand side.
    rhs : expr-like, optional
        The right-hand side. Defaults to 0.
    subdomain : SubDomain, optional
        To restrict the evalaution of the relation to a particular sub-region in the
        computational domain.

    Examples
    --------
    Le may be used to express a relation (e.g. in a Subdomain).

    >>> from devito import Grid, Function, ConditionalDimension, Eq, Operator
    >>> from devito.types import Lt
    >>> grid = Grid(shape=(8, 8))
    >>> g = Function(name='g', shape=grid.shape, dimensions=grid.dimensions)
    >>> x, y = grid.dimensions
    >>> cond = Lt(g, 1)
    >>> ci = ConditionalDimension(name='ci', parent=y, condition=cond)
    >>> f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
    >>> op = Operator(Eq(g, f))
    """

    def __init__(self, lhs, rhs=0, subdomain=None, **kwargs):
        obj = sympy.Lt.__new__(self, lhs, rhs, **kwargs)
        obj._subdomain = subdomain
        return obj

    @property
    def subdomain(self):
        """The SubDomain in which the Lt is defined."""
        return self._subdomain

class Ge(sympy.Ge):
    """
    A greater-than (">=") relation between two objects, the left-hand side and the
    right-hand side. Can be used to build conditionals but not directly to
    construct an Operator.

    The left-hand side may be a Function or a SparseFunction. The right-hand
    side may be any arbitrary expressions with numbers, Dimensions, Constants,
    Functions and SparseFunctions as operands.

    Parameters
    ----------
    lhs : Function or SparseFunction
        The left-hand side.
    rhs : expr-like, optional
        The right-hand side. Defaults to 0.
    subdomain : SubDomain, optional
        To restrict the evalaution of the relation to a particular sub-region in the
        computational domain.

    Examples
    --------
    Le may be used to express a relation (e.g. in a Subdomain).

    >>> from devito import Grid, Function, ConditionalDimension, Eq, Operator
    >>> from devito.types import Ge
    >>> grid = Grid(shape=(8, 8))
    >>> g = Function(name='g', shape=grid.shape, dimensions=grid.dimensions)
    >>> x, y = grid.dimensions
    >>> cond = Ge(g, 1)
    >>> ci = ConditionalDimension(name='ci', parent=y, condition=cond)
    >>> f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
    >>> op = Operator(Eq(g, f))
    """

    def __init__(self, lhs, rhs=0, subdomain=None, **kwargs):
        obj = sympy.Ge.__new__(self, lhs, rhs, **kwargs)
        obj._subdomain = subdomain
        return obj

    @property
    def subdomain(self):
        """The SubDomain in which the Ge is defined."""
        return self._subdomain

class Gt(sympy.Gt):
    """
    A strict-greater-than (">") relation between two objects, the left-hand side and the
    right-hand side. Can be used to build conditionals but not directly to
    construct an Operator.

    The left-hand side may be a Function or a SparseFunction. The right-hand
    side may be any arbitrary expressions with numbers, Dimensions, Constants,
    Functions and SparseFunctions as operands.

    Parameters
    ----------
    lhs : Function or SparseFunction
        The left-hand side.
    rhs : expr-like, optional
        The right-hand side. Defaults to 0.
    subdomain : SubDomain, optional
        To restrict the evalaution of the relation to a particular sub-region in the
        computational domain.

    Examples
    --------
    Gt may be used to express a relation.
    >>> from devito import Grid, Function, ConditionalDimension, Eq, Operator
    >>> from devito.types import Gt
    >>> grid = Grid(shape=(8, 8))
    >>> g = Function(name='g', shape=grid.shape, dimensions=grid.dimensions)
    >>> x, y = grid.dimensions
    >>> cond = Gt(g, 1)
    >>> ci = ConditionalDimension(name='ci', parent=y, condition=cond)
    >>> f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
    >>> op = Operator(Eq(g, f))
    """

    def __init__(self, lhs, rhs=0, subdomain=None, **kwargs):
        obj = sympy.Gt.__new__(self, lhs, rhs, **kwargs)
        obj._subdomain = subdomain
        return obj

    @property
    def subdomain(self):
        """The SubDomain in which the Gt is defined."""
        return self._subdomain
<commit_msg>cd: Address comments, improve docs, remove hacky-API test<commit_after>"""User API to specify relationals."""

import sympy

__all__ = ['Le', 'Lt', 'Ge', 'Gt']

class Le(sympy.Le):
    """
    A less-than or equal ("<=") relation between two objects, the left-hand side and the
    right-hand side. Can be used to build conditionals but not directly to
    construct an Operator.

    The left-hand side may be a Function or a SparseFunction. The right-hand
    side may be any arbitrary expressions with numbers, Dimensions, Constants,
    Functions and SparseFunctions as operands.

    Parameters
    ----------
    lhs : Function or SparseFunction
        The left-hand side.
    rhs : expr-like, optional
        The right-hand side. Defaults to 0.
    subdomain : SubDomain, optional
        To restrict the evalaution of the relation to a particular sub-region in the
        computational domain.

    Examples
    --------
    Le may be used to express a relation (e.g. in a Subdomain).

    >>> from devito import Grid, Function, ConditionalDimension, Eq, Operator
    >>> from devito.types import Le
    >>> grid = Grid(shape=(8, 8))
    >>> g = Function(name='g', shape=grid.shape, dimensions=grid.dimensions)
    >>> x, y = grid.dimensions
    >>> cond = Le(g, 1)
    >>> ci = ConditionalDimension(name='ci', parent=y, condition=cond)
    >>> f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
    >>> op = Operator(Eq(g, f))
    """

    def __init__(self, lhs, rhs=0, subdomain=None, **kwargs):
        obj = sympy.Le.__new__(self, lhs, rhs, **kwargs)
        obj._subdomain = subdomain
        return obj

    @property
    def subdomain(self):
        """The SubDomain in which the Le is defined."""
        return self._subdomain

class Lt(sympy.Lt):
    """
    A strict-less-than ("<") relation between two objects, the left-hand side and the
    right-hand side. Can be used to build conditionals but not directly to
    construct an Operator.

    The left-hand side may be a Function or a SparseFunction. The right-hand
    side may be any arbitrary expressions with numbers, Dimensions, Constants,
    Functions and SparseFunctions as operands.

    Parameters
    ----------
    lhs : Function or SparseFunction
        The left-hand side.
    rhs : expr-like, optional
        The right-hand side. Defaults to 0.
    subdomain : SubDomain, optional
        To restrict the evalaution of the relation to a particular sub-region in the
        computational domain.

    Examples
    --------
    Lt may be used to express a relation (e.g. in a Subdomain).

    >>> from devito import Grid, Function, ConditionalDimension, Eq, Operator
    >>> from devito.types import Lt
    >>> grid = Grid(shape=(8, 8))
    >>> g = Function(name='g', shape=grid.shape, dimensions=grid.dimensions)
    >>> x, y = grid.dimensions
    >>> cond = Lt(g, 1)
    >>> ci = ConditionalDimension(name='ci', parent=y, condition=cond)
    >>> f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
    >>> op = Operator(Eq(g, f))
    """

    def __init__(self, lhs, rhs=0, subdomain=None, **kwargs):
        obj = sympy.Lt.__new__(self, lhs, rhs, **kwargs)
        obj._subdomain = subdomain
        return obj

    @property
    def subdomain(self):
        """The SubDomain in which the Lt is defined."""
        return self._subdomain

class Ge(sympy.Ge):
    """
    A greater-than (">=") relation between two objects, the left-hand side and the
    right-hand side. Can be used to build conditionals but not directly to
    construct an Operator.

    The left-hand side may be a Function or a SparseFunction. The right-hand
    side may be any arbitrary expressions with numbers, Dimensions, Constants,
    Functions and SparseFunctions as operands.

    Parameters
    ----------
    lhs : Function or SparseFunction
        The left-hand side.
    rhs : expr-like, optional
        The right-hand side. Defaults to 0.
    subdomain : SubDomain, optional
        To restrict the evalaution of the relation to a particular sub-region in the
        computational domain.

    Examples
    --------
    Le may be used to express a relation (e.g. in a Subdomain).

    >>> from devito import Grid, Function, ConditionalDimension, Eq, Operator
    >>> from devito.types import Ge
    >>> grid = Grid(shape=(8, 8))
    >>> g = Function(name='g', shape=grid.shape, dimensions=grid.dimensions)
    >>> x, y = grid.dimensions
    >>> cond = Ge(g, 1)
    >>> ci = ConditionalDimension(name='ci', parent=y, condition=cond)
    >>> f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
    >>> op = Operator(Eq(g, f))
    """

    def __init__(self, lhs, rhs=0, subdomain=None, **kwargs):
        obj = sympy.Ge.__new__(self, lhs, rhs, **kwargs)
        obj._subdomain = subdomain
        return obj

    @property
    def subdomain(self):
        """The SubDomain in which the Ge is defined."""
        return self._subdomain

class Gt(sympy.Gt):
    """
    A strict-greater-than (">") relation between two objects, the left-hand side and the
    right-hand side. Can be used to build conditionals but not directly to
    construct an Operator.

    The left-hand side may be a Function or a SparseFunction. The right-hand
    side may be any arbitrary expressions with numbers, Dimensions, Constants,
    Functions and SparseFunctions as operands.

    Parameters
    ----------
    lhs : Function or SparseFunction
        The left-hand side.
    rhs : expr-like, optional
        The right-hand side. Defaults to 0.
    subdomain : SubDomain, optional
        To restrict the evalaution of the relation to a particular sub-region in the
        computational domain.

    Examples
    --------
    Gt may be used to express a relation.
    >>> from devito import Grid, Function, ConditionalDimension, Eq, Operator
    >>> from devito.types import Gt
    >>> grid = Grid(shape=(8, 8))
    >>> g = Function(name='g', shape=grid.shape, dimensions=grid.dimensions)
    >>> x, y = grid.dimensions
    >>> cond = Gt(g, 1)
    >>> ci = ConditionalDimension(name='ci', parent=y, condition=cond)
    >>> f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
    >>> op = Operator(Eq(g, f))
    """

    def __init__(self, lhs, rhs=0, subdomain=None, **kwargs):
        obj = sympy.Gt.__new__(self, lhs, rhs, **kwargs)
        obj._subdomain = subdomain
        return obj

    @property
    def subdomain(self):
        """The SubDomain in which the Gt is defined."""
        return self._subdomain
<|endoftext|>