<commit_before>
#################
# CRS CLASSES
#################

# first classes for each crs element, from the proj4 common paramter listing
# https://trac.osgeo.org/proj/wiki/GenParms

# note also that in wkt the names of "PROJECTION", "DATUM", and "SPHEROID"
# matter and seem to be interpreted and computed upon, ie they carry meaning
# that is commonly understood by programs, so is not explicit in the crs specification.

# the paramters below simply modify certain aspects of the proj/datum/spheroids

# note that the names in wkt of "PROJCS" and "GEOGCS" seem to be purely
# for identifying and branding and can be changed at will.
# they do however act as shortcuts, so that proj4 can use +init=... to load
# everything automatically

# some of these names imply certain combinations of datum and spheroid and paramters.
# so in proj4 one simply needs to give that name, but in wkt one needs to spell it all out.

from . import directions

################

def find(paramname, crstype, strict=False):
    if not strict:
        paramname = paramname.lower()
    for itemname,item in globals().items():
        if itemname.startswith("_"):
            continue
        try:
            if hasattr(item, crstype):
                itemname = getattr(item, crstype)
                if not strict:
                    itemname = itemname.lower()
                if paramname == itemname:
                    return item
        except:
            pass
    else:
        return None

##################
# Ellipsoid parameters

##+a         Semimajor radius of the ellipsoid axis
class SemiMajorRadius:
    proj4 = "+a"
    def __init__(self, value):
        self.value = value

    def to_proj4(self):
        return "%s=%s" % (self.proj4, self.value)

    def to_esri_wkt(self):
        return str(self.value)

    def to_ogc_wkt(self):
        return str(self.value)

##+b         Semiminor radius of the ellipsoid axis
class SemiMinorRadius:
    proj4 = "+b"
    def __init__(self, value):
        self.value = value

    def to_proj4(self):
        return "%s=%s" % (self.proj4, self.value)

    def to_esri_wkt(self):
        return str(self.value)

    def to_ogc_wkt(self):
        return str(self.value)

##+f         Flattening of the ellipsoid axis
class Flattening:
    proj4 = "+f"
    def __init__(self, value):
        self.value = value

    def to_proj4(self):
        return "%s=%s" % (self.proj4, self.value)

    def to_esri_wkt(self):
        return str(self.value)

    def to_ogc_wkt(self):
        return str(self.value)

##+rf         Inverse flattening of the ellipsoid axis
class InverseFlattening:
    proj4 = "+rf"
    def __init__(self, value):
        self.value = value

    def to_proj4(self):
        return "%s=%s" % (self.proj4, self.value)

    def to_esri_wkt(self):
        return str(self.value)

    def to_ogc_wkt(self):
        return str(self.value)

#################
# Other parameters

##+alpha     ? Used with Oblique Mercator and possibly a few others
class Azimuth:
    proj4 = "+alpha"
    esri_wkt = "azimuth"
    ogc_wkt = "azimuth"
    geotiff = "AzimuthAngle"

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

    def to_proj4(self):
        return "+alpha=%s" % self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["Azimuth",%s]' % self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Azimuth",%s]' % self.value

##+k         Scaling factor (old name)
##+k_0       Scaling factor (new name)
class ScalingFactor:
    proj4 = "+k"
    ogc_wkt = "scale_factor"

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

    def to_proj4(self):
        return "+k_0=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["scale_factor", %s]' %self.value

    def to_esri_wkt(self):
        # REALLY??
        raise Exception("Parameter %r not supported by ESRI WKT" % self)

    def to_geotiff(self):
        pass
        #return "ScaleAtNatOrigin" # or ScaleAtCenter?

##+lat_0     Latitude of origin
class LatitudeOrigin:
    proj4 = "+lat_0"
    ogc_wkt = "latitude_of_origin"
    esri_wkt = "Latitude_Of_Origin"

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

    def to_proj4(self):
        return "+lat_0=%s" %self.value

    def to_ogc_wkt(self):
        # SAME AS LATITUDE OF CENTER???
        return 'PARAMETER["latitude_of_origin", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Latitude_Of_Origin", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "ProjCenterLat"

##+lat_1     Latitude of first standard parallel
class LatitudeFirstStndParallel:
    proj4 = "+lat_1"
    ogc_wkt = "standard_parallel_1"
    esri_wkt = "Standard_Parallel_1"

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

    def to_proj4(self):
        return "+lat_1=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["standard_parallel_1", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Standard_Parallel_1", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "StdParallel1"

##+lat_2     Latitude of second standard parallel
class LatitudeSecondStndParallel:
    proj4 = "+lat_2"
    ogc_wkt = "standard_parallel_2"
    esri_wkt = "Standard_Parallel_2"

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

    def to_proj4(self):
        return "+lat_2=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["standard_parallel_2", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Standard_Parallel_2", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "StdParallel2"

##+lat_ts    Latitude of true scale
class LatitudeTrueScale:
    proj4 = "lat_ts"
    ogc_wkt = "Standard_Parallel_1"
    esri_wkt = "Standard_Parallel_1"

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

    def to_proj4(self):
        return "+lat_ts=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["Standard_Parallel_1", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Standard_Parallel_1", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "ProjStdParallel1"

##+lon_0     Central meridian
class CentralMeridian:
    proj4 = "+lon_0"
    ogc_wkt = "Central_Meridian"
    esri_wkt = "Central_Meridian"

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

    def to_proj4(self):
        return "+lon_0=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["Central_Meridian", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Central_Meridian", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "ProjCenterLong"

##+lonc      ? Longitude used with Oblique Mercator and possibly a few others
class LongitudeCenter:
    proj4 = "+lonc"
    ogc_wkt = "Longitude_Of_Center"
    esri_wkt = "Longitude_Of_Center"

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

    def to_proj4(self):
        return "+lonc=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["Longitude_Of_Center", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Longitude_Of_Center", %s]' %self.value

##+lon_wrap  Center longitude to use for wrapping (see below)

##+over      Allow longitude output outside -180 to 180 range, disables wrapping (see below)

##+pm        Alternate prime meridian (typically a city name, see below)
class PrimeMeridian:
    proj4 = "+pm"
    ogc_wkt = "PRIMEM"
    esri_wkt = "PRIMEM"

    cities = {
        'greenwich': 0,
        'lisbon': -9.131906111111112,
        'paris': 2.337229166666667,
        'bogota': -74.08091666666667,
        'madrid': -3.687911111111111,
        'rome': 12.452333333333332,
        'bern': 7.439583333333333,
        'jakarta': 106.80771944444444,
        'ferro': -17.666666666666668,
        'brussels': 4.3679749999999995,
        'stockholm': 18.05827777777778,
        'athens': 23.716337499999998,
        'oslo': 10.722916666666666,
        }

    # based on http://proj.maptools.org/gen_parms.html
##       greenwich 0dE
##          lisbon 9d07'54.862"W
##           paris 2d20'14.025"E
##          bogota 74d04'51.3"E
##          madrid 3d41'16.48"W
##            rome 12d27'8.4"E
##            bern 7d26'22.5"E
##         jakarta 106d48'27.79"E
##           ferro 17d40'W
##        brussels 4d22'4.71"E
##       stockholm 18d3'29.8"E
##          athens 23d42'58.815"E
##            oslo 10d43'22.5"E

    def __init__(self, value):
        """
        The prime meridian coordinate (or city name), relative to greenwhich, where the
        longitude is considered to be 0.

        Arguments:

        - **value**: Longitude value relative to Greenwich, or the name of one of these cities: {}.
        """.format(', '.join(self.cities.keys()))
        self.value = value

    def get_value(self):
        value = self.value
        try:
            value = float(value)
        except:
            value = value.lower()
            if value not in self.cities:
                raise Exception("Prime meridian value {} must be a number or the name of one of these cities: {}".format(value, ', '.join(self.cities.keys())))
            value = self.cities[value]
        value = int(value) if value.is_integer() else value
        return value

    def to_proj4(self):
        return "+pm=%s" %self.get_value()

    def to_ogc_wkt(self):
        return 'PRIMEM["Greenwich", %s]' %self.get_value()

    def to_esri_wkt(self):
        return 'PRIMEM["Greenwich", %s]' %self.get_value()

##+zone     UTM zone

##+south     Denotes southern hemisphere UTM zone

##+towgs84   3 or 7 term datum transform parameters (see below)
class DatumShift:
    proj4 = "+towgs84"
    ogc_wkt = "TOWGS84"

    def __init__(self, value):
        """
        The WGS84 Datum shift parameter.

        Args:

        - **value**: A list of 3 or 7 term datum transform parameters.
        """
        self.value = value

    def to_proj4(self):
        return "+towgs84=%s" %",".join((str(val) for val in self.value))

    def to_ogc_wkt(self):
        return "TOWGS84[%s]" %",".join((str(val) for val in self.value))

    def to_esri_wkt(self):
        raise Exception("Parameter %r not supported by ESRI WKT" % self)

##+x_0       False easting
class FalseEasting:
    proj4 = "+x_0"
    esri_wkt = "False_Easting"
    ogc_wkt = "false_easting"
    geotiff = "FalseEasting"

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

    def to_proj4(self):
        return "+x_0=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["false_easting", %s]' % self.value

    def to_esri_wkt(self):
        return 'PARAMETER["False_Easting", %s]' % self.value

##+y_0       False northing
class FalseNorthing:
    proj4 = "+y_0"
    esri_wkt = "False_Northing"
    ogc_wkt = "false_northing"
    geotiff = "FalseNorthing"

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

    def to_proj4(self):
        return "+y_0=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["false_northing", %s]' % self.value

    def to_esri_wkt(self):
        return 'PARAMETER["False_Northing", %s]' % self.value

##+h       Satellite height
class SatelliteHeight:
    proj4 = "+h"
    ogc_wkt = "satellite_height"
    esri_wkt = "satellite_height"

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

    def to_proj4(self):
        return "+h=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["satellite_height", %s]' % self.value

    def to_esri_wkt(self):
        return 'PARAMETER["satellite_height", %s]' % self.value

##+tilt     Tilt angle
class TiltAngle:
    proj4 = "+tilt"

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

    def to_proj4(self):
        return "+tilt=%s" %self.value

    def to_ogc_wkt(self):
        raise Exception("Parameter not supported by OGC WKT")

    def to_esri_wkt(self):
        raise Exception("Parameter not supported by ESRI WKT")

<commit_msg>Added scale factor as possibility for ESRI WKT<commit_after>
#################
# CRS CLASSES
#################

# first classes for each crs element, from the proj4 common paramter listing
# https://trac.osgeo.org/proj/wiki/GenParms

# note also that in wkt the names of "PROJECTION", "DATUM", and "SPHEROID"
# matter and seem to be interpreted and computed upon, ie they carry meaning
# that is commonly understood by programs, so is not explicit in the crs specification.

# the paramters below simply modify certain aspects of the proj/datum/spheroids

# note that the names in wkt of "PROJCS" and "GEOGCS" seem to be purely
# for identifying and branding and can be changed at will.
# they do however act as shortcuts, so that proj4 can use +init=... to load
# everything automatically

# some of these names imply certain combinations of datum and spheroid and paramters.
# so in proj4 one simply needs to give that name, but in wkt one needs to spell it all out.

from . import directions

################

def find(paramname, crstype, strict=False):
    if not strict:
        paramname = paramname.lower()
    for itemname,item in globals().items():
        if itemname.startswith("_"):
            continue
        try:
            if hasattr(item, crstype):
                itemname = getattr(item, crstype)
                if not strict:
                    itemname = itemname.lower()
                if paramname == itemname:
                    return item
        except:
            pass
    else:
        return None

##################
# Ellipsoid parameters

##+a         Semimajor radius of the ellipsoid axis
class SemiMajorRadius:
    proj4 = "+a"
    def __init__(self, value):
        self.value = value

    def to_proj4(self):
        return "%s=%s" % (self.proj4, self.value)

    def to_esri_wkt(self):
        return str(self.value)

    def to_ogc_wkt(self):
        return str(self.value)

##+b         Semiminor radius of the ellipsoid axis
class SemiMinorRadius:
    proj4 = "+b"
    def __init__(self, value):
        self.value = value

    def to_proj4(self):
        return "%s=%s" % (self.proj4, self.value)

    def to_esri_wkt(self):
        return str(self.value)

    def to_ogc_wkt(self):
        return str(self.value)

##+f         Flattening of the ellipsoid axis
class Flattening:
    proj4 = "+f"
    def __init__(self, value):
        self.value = value

    def to_proj4(self):
        return "%s=%s" % (self.proj4, self.value)

    def to_esri_wkt(self):
        return str(self.value)

    def to_ogc_wkt(self):
        return str(self.value)

##+rf         Inverse flattening of the ellipsoid axis
class InverseFlattening:
    proj4 = "+rf"
    def __init__(self, value):
        self.value = value

    def to_proj4(self):
        return "%s=%s" % (self.proj4, self.value)

    def to_esri_wkt(self):
        return str(self.value)

    def to_ogc_wkt(self):
        return str(self.value)

#################
# Other parameters

##+alpha     ? Used with Oblique Mercator and possibly a few others
class Azimuth:
    proj4 = "+alpha"
    esri_wkt = "azimuth"
    ogc_wkt = "azimuth"
    geotiff = "AzimuthAngle"

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

    def to_proj4(self):
        return "+alpha=%s" % self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["Azimuth",%s]' % self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Azimuth",%s]' % self.value

##+k         Scaling factor (old name)
##+k_0       Scaling factor (new name)
class ScalingFactor:
    proj4 = "+k"
    esri_wkt = "Scale_Factor"
    ogc_wkt = "scale_factor"

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

    def to_proj4(self):
        return "+k_0=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["scale_factor", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Scale_Factor", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "ScaleAtNatOrigin" # or ScaleAtCenter?

##+lat_0     Latitude of origin
class LatitudeOrigin:
    proj4 = "+lat_0"
    ogc_wkt = "latitude_of_origin"
    esri_wkt = "Latitude_Of_Origin"

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

    def to_proj4(self):
        return "+lat_0=%s" %self.value

    def to_ogc_wkt(self):
        # SAME AS LATITUDE OF CENTER???
        return 'PARAMETER["latitude_of_origin", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Latitude_Of_Origin", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "ProjCenterLat"

##+lat_1     Latitude of first standard parallel
class LatitudeFirstStndParallel:
    proj4 = "+lat_1"
    ogc_wkt = "standard_parallel_1"
    esri_wkt = "Standard_Parallel_1"

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

    def to_proj4(self):
        return "+lat_1=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["standard_parallel_1", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Standard_Parallel_1", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "StdParallel1"

##+lat_2     Latitude of second standard parallel
class LatitudeSecondStndParallel:
    proj4 = "+lat_2"
    ogc_wkt = "standard_parallel_2"
    esri_wkt = "Standard_Parallel_2"

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

    def to_proj4(self):
        return "+lat_2=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["standard_parallel_2", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Standard_Parallel_2", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "StdParallel2"

##+lat_ts    Latitude of true scale
class LatitudeTrueScale:
    proj4 = "lat_ts"
    ogc_wkt = "Standard_Parallel_1"
    esri_wkt = "Standard_Parallel_1"

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

    def to_proj4(self):
        return "+lat_ts=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["Standard_Parallel_1", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Standard_Parallel_1", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "ProjStdParallel1"

##+lon_0     Central meridian
class CentralMeridian:
    proj4 = "+lon_0"
    ogc_wkt = "Central_Meridian"
    esri_wkt = "Central_Meridian"

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

    def to_proj4(self):
        return "+lon_0=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["Central_Meridian", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Central_Meridian", %s]' %self.value

    def to_geotiff(self):
        pass
        #return "ProjCenterLong"

##+lonc      ? Longitude used with Oblique Mercator and possibly a few others
class LongitudeCenter:
    proj4 = "+lonc"
    ogc_wkt = "Longitude_Of_Center"
    esri_wkt = "Longitude_Of_Center"

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

    def to_proj4(self):
        return "+lonc=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["Longitude_Of_Center", %s]' %self.value

    def to_esri_wkt(self):
        return 'PARAMETER["Longitude_Of_Center", %s]' %self.value

##+lon_wrap  Center longitude to use for wrapping (see below)

##+over      Allow longitude output outside -180 to 180 range, disables wrapping (see below)

##+pm        Alternate prime meridian (typically a city name, see below)
class PrimeMeridian:
    proj4 = "+pm"
    ogc_wkt = "PRIMEM"
    esri_wkt = "PRIMEM"

    cities = {
        'greenwich': 0,
        'lisbon': -9.131906111111112,
        'paris': 2.337229166666667,
        'bogota': -74.08091666666667,
        'madrid': -3.687911111111111,
        'rome': 12.452333333333332,
        'bern': 7.439583333333333,
        'jakarta': 106.80771944444444,
        'ferro': -17.666666666666668,
        'brussels': 4.3679749999999995,
        'stockholm': 18.05827777777778,
        'athens': 23.716337499999998,
        'oslo': 10.722916666666666,
        }

    # based on http://proj.maptools.org/gen_parms.html
##       greenwich 0dE
##          lisbon 9d07'54.862"W
##           paris 2d20'14.025"E
##          bogota 74d04'51.3"E
##          madrid 3d41'16.48"W
##            rome 12d27'8.4"E
##            bern 7d26'22.5"E
##         jakarta 106d48'27.79"E
##           ferro 17d40'W
##        brussels 4d22'4.71"E
##       stockholm 18d3'29.8"E
##          athens 23d42'58.815"E
##            oslo 10d43'22.5"E

    def __init__(self, value):
        """
        The prime meridian coordinate (or city name), relative to greenwhich, where the
        longitude is considered to be 0.

        Arguments:

        - **value**: Longitude value relative to Greenwich, or the name of one of these cities: {}.
        """.format(', '.join(self.cities.keys()))
        self.value = value

    def get_value(self):
        value = self.value
        try:
            value = float(value)
        except:
            value = value.lower()
            if value not in self.cities:
                raise Exception("Prime meridian value {} must be a number or the name of one of these cities: {}".format(value, ', '.join(self.cities.keys())))
            value = self.cities[value]
        value = int(value) if value.is_integer() else value
        return value

    def to_proj4(self):
        return "+pm=%s" %self.get_value()

    def to_ogc_wkt(self):
        return 'PRIMEM["Greenwich", %s]' %self.get_value()

    def to_esri_wkt(self):
        return 'PRIMEM["Greenwich", %s]' %self.get_value()

##+zone     UTM zone

##+south     Denotes southern hemisphere UTM zone

##+towgs84   3 or 7 term datum transform parameters (see below)
class DatumShift:
    proj4 = "+towgs84"
    ogc_wkt = "TOWGS84"

    def __init__(self, value):
        """
        The WGS84 Datum shift parameter.

        Args:

        - **value**: A list of 3 or 7 term datum transform parameters.
        """
        self.value = value

    def to_proj4(self):
        return "+towgs84=%s" %",".join((str(val) for val in self.value))

    def to_ogc_wkt(self):
        return "TOWGS84[%s]" %",".join((str(val) for val in self.value))

    def to_esri_wkt(self):
        raise Exception("Parameter %r not supported by ESRI WKT" % self)

##+x_0       False easting
class FalseEasting:
    proj4 = "+x_0"
    esri_wkt = "False_Easting"
    ogc_wkt = "false_easting"
    geotiff = "FalseEasting"

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

    def to_proj4(self):
        return "+x_0=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["false_easting", %s]' % self.value

    def to_esri_wkt(self):
        return 'PARAMETER["False_Easting", %s]' % self.value

##+y_0       False northing
class FalseNorthing:
    proj4 = "+y_0"
    esri_wkt = "False_Northing"
    ogc_wkt = "false_northing"
    geotiff = "FalseNorthing"

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

    def to_proj4(self):
        return "+y_0=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["false_northing", %s]' % self.value

    def to_esri_wkt(self):
        return 'PARAMETER["False_Northing", %s]' % self.value

##+h       Satellite height
class SatelliteHeight:
    proj4 = "+h"
    ogc_wkt = "satellite_height"
    esri_wkt = "satellite_height"

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

    def to_proj4(self):
        return "+h=%s" %self.value

    def to_ogc_wkt(self):
        return 'PARAMETER["satellite_height", %s]' % self.value

    def to_esri_wkt(self):
        return 'PARAMETER["satellite_height", %s]' % self.value

##+tilt     Tilt angle
class TiltAngle:
    proj4 = "+tilt"

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

    def to_proj4(self):
        return "+tilt=%s" %self.value

    def to_ogc_wkt(self):
        raise Exception("Parameter not supported by OGC WKT")

    def to_esri_wkt(self):
        raise Exception("Parameter not supported by ESRI WKT")

<|endoftext|>


<commit_before>import os
import time
import shutil
import json
import sys

import config
import dockerapi as docker
import util
import network

if util.on_macos:
    import OSes.macos as OS

elif util.on_wsl:
    import OSes.windows_wsl2 as OS

elif util.on_linux:
    if config.NAME == 'Ubuntu':
        import OSes.ubuntu as OS
    elif config.NAME.lower() == 'linux mint':
        import OSes.mint as OS

RESOLVCONF = '/etc/resolv.conf'
RESOLVCONF_HEAD = '/etc/resolvconf/resolv.conf.d/head'
RESOLVCONF_TAIL = '/etc/resolvconf/resolv.conf.d/tail'

def update_cache():
    util.write_cache('tld', config.TOP_LEVEL_DOMAIN)
    util.write_cache('tag', config.DOCKER_CONTAINER_TAG)
    util.write_cache('name', config.DOCKER_CONTAINER_NAME)

def update_resolvconf():
    if not hasattr(OS, 'DISABLE_MAIN_RESOLVCONF_ROUTINE'):
        if not os.path.exists(RESOLVCONF):
            content = '# Generated by docker-dns\n'
            open(RESOLVCONF, 'w').write(content)

        dns = docker.NETWORK_GATEWAY
        try:
            dns = OS.DNS
        except Exception:
            # no DNS
            pass

        name_servers = f'nameserver {dns} #@docker-dns\n' \
                       f'nameserver 1.1.1.1 #@docker-dns\n' \
                       f'nameserver 8.8.8.8 #@docker-dns\n'
        options = 'options timeout:1 #@docker-dns\n'
        if OS.FLAVOR == 'ubuntu' and config.OS_VERSION >= 18 * 1000:
            open(RESOLVCONF_HEAD, 'a').write(name_servers)
            open(RESOLVCONF_TAIL, 'a').write(options)

        open(RESOLVCONF, 'a').write(f'{name_servers}{options}')

def main(name=config.DOCKER_CONTAINER_NAME, tag=config.DOCKER_CONTAINER_TAG, tld=config.TOP_LEVEL_DOMAIN):
    if not util.is_os_supported(OS.FLAVOR):
        return 1

    if os.path.exists('.cache/INSTALLED'):
        os.unlink('.cache/INSTALLED')

    try:
        update_resolvconf()
    except FileNotFoundError:
        print('package "resolvconf" not installed, please install it with "sudo apt install resolvconf"')
        return 1

    # docker
    docker_conf_file = f"{OS.DOCKER_CONF_FOLDER}/daemon.json"
    if not os.path.exists(docker_conf_file) or os.stat(docker_conf_file).st_size == 0:
        if not os.path.isdir(OS.DOCKER_CONF_FOLDER):
            os.mkdir(OS.DOCKER_CONF_FOLDER)
        shutil.copy2('src/templates/daemon.json', docker_conf_file)

    docker_json = json.loads(open(docker_conf_file, 'r').read())
    docker_json['bip'] = docker.NETWORK_SUBNET
    docker_json['dns'] = list(
        set([docker.NETWORK_GATEWAY] + network.get_dns_servers()))
    with open(docker_conf_file, 'w') as daemon_file:
        daemon_file.write(json.dumps(docker_json, indent=4, sort_keys=True))

    if docker.check_exists(name):
        print("Stopping existing container...")
        docker.purge(name)
        time.sleep(2)

    print(
        f'Building and running container "{tag}:latest"... Please wait')
    docker.build_container(
        name, tag, tld, bind_port_ip=util.on_linux and not util.on_wsl, target=OS.DOCKER_BUILD_TARGET)
    update_cache()

    # TLD domain certificate
    cert_file = f'certs.d/tld/{tld}.cert'
    key_file = f'certs.d/tld/{tld}.key'
    util.generate_certificate(tld, cert_file=cert_file, key_file=key_file)
    shutil.copy2(cert_file, OS.DOCKER_CONF_FOLDER)
    shutil.copy2(key_file, OS.DOCKER_CONF_FOLDER)

    original_arg = sys.argv
    install_status = OS.install(tld)
    if install_status and util.is_tunnel_needed():
        original_arg[1] = 'tunnel'
        original_arg.append('&')
        os.system(' '.join(original_arg))

    open('.cache/INSTALLED', 'w').write('')
    return 0
<commit_msg>refactored dns expression<commit_after>import os
import time
import shutil
import json
import sys

import config
import dockerapi as docker
import util
import network

if util.on_macos:
    import OSes.macos as OS

elif util.on_wsl:
    import OSes.windows_wsl2 as OS

elif util.on_linux:
    if config.NAME == 'Ubuntu':
        import OSes.ubuntu as OS
    elif config.NAME.lower() == 'linux mint':
        import OSes.mint as OS

RESOLVCONF = '/etc/resolv.conf'
RESOLVCONF_HEAD = '/etc/resolvconf/resolv.conf.d/head'
RESOLVCONF_TAIL = '/etc/resolvconf/resolv.conf.d/tail'

def update_cache():
    util.write_cache('tld', config.TOP_LEVEL_DOMAIN)
    util.write_cache('tag', config.DOCKER_CONTAINER_TAG)
    util.write_cache('name', config.DOCKER_CONTAINER_NAME)

def update_resolvconf():
    if not hasattr(OS, 'DISABLE_MAIN_RESOLVCONF_ROUTINE'):
        if not os.path.exists(RESOLVCONF):
            content = '# Generated by docker-dns\n'
            open(RESOLVCONF, 'w').write(content)

        dns = docker.NETWORK_GATEWAY
        try:
            dns = OS.DNS
        except Exception:
            # no DNS
            pass

        name_servers = f'nameserver {dns} #@docker-dns\n' \
                       f'nameserver 1.1.1.1 #@docker-dns\n' \
                       f'nameserver 8.8.8.8 #@docker-dns\n'
        options = 'options timeout:1 #@docker-dns\n'
        if OS.FLAVOR == 'ubuntu' and config.OS_VERSION >= 18 * 1000:
            open(RESOLVCONF_HEAD, 'a').write(name_servers)
            open(RESOLVCONF_TAIL, 'a').write(options)

        open(RESOLVCONF, 'a').write(f'{name_servers}{options}')

def main(name=config.DOCKER_CONTAINER_NAME, tag=config.DOCKER_CONTAINER_TAG, tld=config.TOP_LEVEL_DOMAIN):
    if not util.is_os_supported(OS.FLAVOR):
        return 1

    if os.path.exists('.cache/INSTALLED'):
        os.unlink('.cache/INSTALLED')

    try:
        update_resolvconf()
    except FileNotFoundError:
        print('package "resolvconf" not installed, please install it with "sudo apt install resolvconf"')
        return 1

    # docker
    docker_conf_file = f"{OS.DOCKER_CONF_FOLDER}/daemon.json"
    if not os.path.exists(docker_conf_file) or os.stat(docker_conf_file).st_size == 0:
        if not os.path.isdir(OS.DOCKER_CONF_FOLDER):
            os.mkdir(OS.DOCKER_CONF_FOLDER)
        shutil.copy2('src/templates/daemon.json', docker_conf_file)

    docker_json = json.loads(open(docker_conf_file, 'r').read())
    docker_json['bip'] = docker.NETWORK_SUBNET
    docker_json['dns'] = [{docker.NETWORK_GATEWAY, *network.get_dns_servers()}]
    with open(docker_conf_file, 'w') as daemon_file:
        daemon_file.write(json.dumps(docker_json, indent=4, sort_keys=True))

    if docker.check_exists(name):
        print("Stopping existing container...")
        docker.purge(name)
        time.sleep(2)

    print(
        f'Building and running container "{tag}:latest"... Please wait')
    docker.build_container(
        name, tag, tld, bind_port_ip=util.on_linux and not util.on_wsl, target=OS.DOCKER_BUILD_TARGET)
    update_cache()

    # TLD domain certificate
    cert_file = f'certs.d/tld/{tld}.cert'
    key_file = f'certs.d/tld/{tld}.key'
    util.generate_certificate(tld, cert_file=cert_file, key_file=key_file)
    shutil.copy2(cert_file, OS.DOCKER_CONF_FOLDER)
    shutil.copy2(key_file, OS.DOCKER_CONF_FOLDER)

    original_arg = sys.argv
    install_status = OS.install(tld)
    if install_status and util.is_tunnel_needed():
        original_arg[1] = 'tunnel'
        original_arg.append('&')
        os.system(' '.join(original_arg))

    open('.cache/INSTALLED', 'w').write('')
    return 0
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Download and extract indices from EDGAR index files.

:copyright: © 2020 by Mac Gaulin
:license: MIT, see LICENSE for more details.
"""

# Stdlib imports
import os
import logging

# 3rd party imports
import pandas as pd

# Module Imports
from pyedgar import config
from pyedgar import utilities
from pyedgar.utilities import edgarweb

# progress logging
try:
    from tqdm import tqdm
except ImportError:

    def tqdm(_iterable, *args, **kwargs):
        return _iterable

class IndexMaker:
    """
    Class that downloads EDGAR to your very own computer.
    Everyone should have a local copy after all!
    I have just enough flexibility here for my computer. It works on Linux, YWindowsMMV.
    """

    edgar_index_args = {"sep": "|", "encoding": "latin-1", "skiprows": [0, 1, 2, 3, 4, 5, 6, 7, 8, 10]}

    _use_requests = False
    _tq = None
    _get_index_cache_path = None

    # May as well share this across instances (instead of setting in __init__)
    _logger = logging.getLogger(__name__)

    def __init__(self, use_tqdm=False, use_requests=False):
        """
        Initialize the index making object.

        Arguments:
            use_tqdm (bool): flag for whether or not to wrap downloads in tqdm for progress monitoring
        """
        self._use_requests = use_requests
        self._get_index_cache_path = config.get_index_cache_path

        if use_tqdm:
            self._tq = tqdm
        else:

            def _tq(_iterable, *args, **kwargs):
                return _iterable

            self._tq = _tq

    def download_indexes(self, start_year=1995, stop_year=None, overwrite=False):
        """Download multiple edgar quarterly index compressed files.

        Args:
            start_year (int, datetime, None): Starting year (or datetime from which we'll extract year/quarter). Default to 1995.
            stop_year (int, datetime, None): Ending year (or datetime from which we'll extract year/quarter). Default to today's year.
            overwrite (bool): Flag for whether to overwrite any existing file (default False).
            use_requests (bool): Flag for whether to use requests or curl (default False == curl).
            overwrite_size_threshold (int): Existing files smaller than this will be re-downloaded.

        Returns:
            tuple: output file path, return code
        """
        _num = len([0 for _ in utilities.iterate_dates(start_year, stop_year, period="quarterly")])

        # The download recursively works as it sounds, but we want progress bar, so do it date by date
        for i_date in self._tq(
            utilities.iterate_dates(start_year, stop_year, period="quarterly"), total=_num, desc="Downloading Indices"
        ):
            edgarweb.download_indexes_recursively(
                i_date, end_date=i_date, overwrite=overwrite, use_requests=self._use_requests
            )

    def extract_indexes(self, start_year=1995, stop_year=None, save_forms=None, download_first=True, overwrite=False):
        if download_first:
            self._logger.info("Downloading the quarterly indices...")
            self.download_indexes(start_year=start_year, stop_year=stop_year, overwrite=overwrite)
            self._logger.info("Done downloading quarterly indices.")

        df = pd.DataFrame()

        _num = len([0 for _ in utilities.iterate_dates(start_year, stop_year, period="quarterly")])

        for i_date in self._tq(
            utilities.iterate_dates(start_year, to_date=stop_year, period="quarterly"),
            total=_num,
            desc="Extracting Indices",
        ):
            idx_cache_file = self._get_index_cache_path(i_date)

            try:
                self._logger.info("\tLoading index for %rQ%r", i_date.year, utilities.get_quarter(i_date))
                dfi = pd.read_csv(idx_cache_file, **self.edgar_index_args)
            except FileNotFoundError:
                self._logger.warning(
                    "No Index cache file at %r (for %rQ%r)", idx_cache_file, i_date.year, utilities.get_quarter(i_date)
                )
            except Exception:
                self._logger.warning(
                    "Reading %r failed at %rQ%r", idx_cache_file, i_date.year, utilities.get_quarter(i_date)
                )
                # File reading didn't work, try overwriting with new download
                continue

            dfi["Accession"] = dfi.Filename.str.slice(start=-24, stop=-4)
            del dfi["Filename"]

            df = pd.concat([df, dfi], copy=False)

            self._logger.info("Added %r(%r) to df(%r)", idx_cache_file, len(dfi), len(df))

        df["Date Filed"] = pd.to_datetime(df["Date Filed"])

        all_forms = df["Form Type"].unique()
        if save_forms is None:
            save_forms = {
                "all": all_forms,
                "10-K": [x for x in all_forms if x[:4] in ("10-K", "10KS")],
                "10-Q": [x for x in all_forms if x[:4] in ("10-Q", "10QS")],
                "DEF14A": [x for x in all_forms if x.endswith("14A")],
                "8-K": ("8-K", "8-K/A"),
            }

        for form, formlist in self._tq(save_forms.items(), total=len(save_forms), desc="Exporting Indices"):
            outpath = os.path.join(config.INDEX_ROOT, "form_{}.{}".format(form, config.INDEX_EXTENSION))

            self._logger.info("Saving %r: %r", outpath, formlist)

            (
                df[df["Form Type"].isin(formlist)]
                .sort_values(["CIK", "Date Filed"])
                .to_csv(outpath, sep=config.INDEX_DELIMITER, index=False)
            )
        self._logger.info("Done extracting indices!")
<commit_msg>Reduce info level logging in index download<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Download and extract indices from EDGAR index files.

:copyright: © 2020 by Mac Gaulin
:license: MIT, see LICENSE for more details.
"""

# Stdlib imports
import os
import logging

# 3rd party imports
import pandas as pd

# Module Imports
from pyedgar import config
from pyedgar import utilities
from pyedgar.utilities import edgarweb

# progress logging
try:
    from tqdm import tqdm
except ImportError:

    def tqdm(_iterable, *args, **kwargs):
        return _iterable

class IndexMaker:
    """
    Class that downloads EDGAR to your very own computer.
    Everyone should have a local copy after all!
    I have just enough flexibility here for my computer. It works on Linux, YWindowsMMV.
    """

    edgar_index_args = {"sep": "|", "encoding": "latin-1", "skiprows": [0, 1, 2, 3, 4, 5, 6, 7, 8, 10]}

    _use_requests = False
    _tq = None
    _get_index_cache_path = None

    # May as well share this across instances (instead of setting in __init__)
    _logger = logging.getLogger(__name__)

    def __init__(self, use_tqdm=False, use_requests=False):
        """
        Initialize the index making object.

        Arguments:
            use_tqdm (bool): flag for whether or not to wrap downloads in tqdm for progress monitoring
        """
        self._use_requests = use_requests
        self._get_index_cache_path = config.get_index_cache_path

        if use_tqdm:
            self._tq = tqdm
        else:

            def _tq(_iterable, *args, **kwargs):
                return _iterable

            self._tq = _tq

    def download_indexes(self, start_year=1995, stop_year=None, overwrite=False):
        """Download multiple edgar quarterly index compressed files.

        Args:
            start_year (int, datetime, None): Starting year (or datetime from which we'll extract year/quarter). Default to 1995.
            stop_year (int, datetime, None): Ending year (or datetime from which we'll extract year/quarter). Default to today's year.
            overwrite (bool): Flag for whether to overwrite any existing file (default False).
            use_requests (bool): Flag for whether to use requests or curl (default False == curl).
            overwrite_size_threshold (int): Existing files smaller than this will be re-downloaded.

        Returns:
            tuple: output file path, return code
        """
        _num = len([0 for _ in utilities.iterate_dates(start_year, stop_year, period="quarterly")])

        # The download recursively works as it sounds, but we want progress bar, so do it date by date
        for i_date in self._tq(
            utilities.iterate_dates(start_year, stop_year, period="quarterly"), total=_num, desc="Downloading Indices"
        ):
            edgarweb.download_indexes_recursively(
                i_date, end_date=i_date, overwrite=overwrite, use_requests=self._use_requests
            )

    def extract_indexes(self, start_year=1995, stop_year=None, save_forms=None, download_first=True, overwrite=False):
        if download_first:
            self._logger.info("Downloading the quarterly indices...")
            self.download_indexes(start_year=start_year, stop_year=stop_year, overwrite=overwrite)
            self._logger.info("Done downloading quarterly indices.")

        df = pd.DataFrame()

        _num = len([0 for _ in utilities.iterate_dates(start_year, stop_year, period="quarterly")])

        for i_date in self._tq(
            utilities.iterate_dates(start_year, to_date=stop_year, period="quarterly"),
            total=_num,
            desc="Extracting Indices",
        ):
            idx_cache_file = self._get_index_cache_path(i_date)

            try:
                self._logger.info("\tLoading index for %rQ%r", i_date.year, utilities.get_quarter(i_date))
                dfi = pd.read_csv(idx_cache_file, **self.edgar_index_args)
            except FileNotFoundError:
                self._logger.warning(
                    "No Index cache file at %r (for %rQ%r)", idx_cache_file, i_date.year, utilities.get_quarter(i_date)
                )
            except Exception:
                self._logger.warning(
                    "Reading %r failed at %rQ%r", idx_cache_file, i_date.year, utilities.get_quarter(i_date)
                )
                # File reading didn't work, try overwriting with new download
                continue

            dfi["Accession"] = dfi.Filename.str.slice(start=-24, stop=-4)
            del dfi["Filename"]

            df = pd.concat([df, dfi], copy=False)

            self._logger.info("Added %r(%r) to df(%r)", idx_cache_file, len(dfi), len(df))

        df["Date Filed"] = pd.to_datetime(df["Date Filed"])

        all_forms = df["Form Type"].unique()
        if save_forms is None:
            save_forms = {
                "all": all_forms,
                "10-K": [x for x in all_forms if x[:4] in ("10-K", "10KS")],
                "10-Q": [x for x in all_forms if x[:4] in ("10-Q", "10QS")],
                "DEF14A": [x for x in all_forms if x.endswith("14A")],
                "8-K": ("8-K", "8-K/A"),
            }

        for form, formlist in self._tq(save_forms.items(), total=len(save_forms), desc="Exporting Indices"):
            outpath = os.path.join(config.INDEX_ROOT, "form_{}.{}".format(form, config.INDEX_EXTENSION))

            self._logger.info("Saving %r to %r", form, outpath)
            self._logger.debug("Saving %s to %s with form-types: %r", form, outpath, formlist)

            (
                df[df["Form Type"].isin(formlist)]
                .sort_values(["CIK", "Date Filed"])
                .to_csv(outpath, sep=config.INDEX_DELIMITER, index=False)
            )
        self._logger.info("Done extracting indices!")
<|endoftext|>


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

r"""Trainer for RL environments.

For now we only support PPO as RL algorithm.

Sample invocation:

TRAIN_BATCH_SIZE=32
python trax/rl_trainer.py \
  --config_file=trax/rl/configs/acrobot.gin \
  --train_batch_size=${TRAIN_BATCH_SIZE} \
  --output_dir=${HOME}/ppo_acrobot \
  --vmodule=*/trax/*=1 \
  --alsologtostderr
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import multiprocessing
import os

from absl import app
from absl import flags
from absl import logging
import gin
import jax
from jax.config import config
from tensor2tensor import envs  # pylint: disable=unused-import
from tensor2tensor.envs import env_problem_utils
from trax import rl  # pylint: disable=unused-import
from trax import trainer_flags  # pylint: disable=unused-import
from trax.rl import envs as rl_envs  # pylint: disable=unused-import
from trax.rl import task as rl_task
from trax.rl import trainers as rl_trainers
from trax.rl import training as light_trainers

FLAGS = flags.FLAGS

# Not just 'train' to avoid a conflict with trax.train in GIN files.
@gin.configurable(blacklist=[
    'output_dir', 'train_batch_size', 'eval_batch_size', 'trajectory_dump_dir'
])
def train_rl(
    output_dir,
    train_batch_size,
    eval_batch_size,
    env_name='Acrobot-v1',
    max_timestep=None,
    clip_rewards=False,
    rendered_env=False,
    resize=False,
    resize_dims=(105, 80),
    trainer_class=rl_trainers.PPO,
    n_epochs=10000,
    trajectory_dump_dir=None,
    num_actions=None,
    light_rl=False,
    light_rl_trainer=light_trainers.RLTrainer,
):
  """Train the RL agent.

  Args:
    output_dir: Output directory.
    train_batch_size: Number of parallel environments to use for training.
    eval_batch_size: Number of parallel environments to use for evaluation.
    env_name: Name of the environment.
    max_timestep: Int or None, the maximum number of timesteps in a trajectory.
      The environment is wrapped in a TimeLimit wrapper.
    clip_rewards: Whether to clip and discretize the rewards.
    rendered_env: Whether the environment has visual input. If so, a
      RenderedEnvProblem will be used.
    resize: whether to do resize or not
    resize_dims: Pair (height, width), dimensions to resize the visual
      observations to.
    trainer_class: RLTrainer class to use.
    n_epochs: Number epochs to run the training for.
    trajectory_dump_dir: Directory to dump trajectories to.
    num_actions: None unless one wants to use the discretization wrapper. Then
      num_actions specifies the number of discrete actions.
    light_rl: whether to use the light RL setting (experimental).
    light_rl_trainer: whichh light RL trainer to use (experimental).
  """

  if FLAGS.jax_debug_nans:
    config.update('jax_debug_nans', True)

  if FLAGS.use_tpu:
    config.update('jax_platform_name', 'tpu')
  else:
    config.update('jax_platform_name', 'gpu')

  if light_rl:
    task = rl_task.RLTask()
    trainer = light_rl_trainer(task=task, output_dir=output_dir)
    trainer.run(n_epochs)
    return

  # TODO(pkozakowski): Find a better way to determine this.
  train_env_kwargs = {}
  eval_env_kwargs = {}
  if 'OnlineTuneEnv' in env_name:
    envs_output_dir = FLAGS.envs_output_dir or os.path.join(output_dir, 'envs')
    train_env_output_dir = os.path.join(envs_output_dir, 'train')
    eval_env_output_dir = os.path.join(envs_output_dir, 'eval')
    train_env_kwargs = {'output_dir': train_env_output_dir}
    eval_env_kwargs = {'output_dir': eval_env_output_dir}

  parallelism = multiprocessing.cpu_count() if FLAGS.parallelize_envs else 1

  logging.info('Num discretized actions %s', num_actions)
  logging.info('Resize %d', resize)

  train_env = env_problem_utils.make_env(
      batch_size=train_batch_size,
      env_problem_name=env_name,
      rendered_env=rendered_env,
      resize=resize,
      resize_dims=resize_dims,
      max_timestep=max_timestep,
      clip_rewards=clip_rewards,
      parallelism=parallelism,
      use_tpu=FLAGS.use_tpu,
      num_actions=num_actions,
      **train_env_kwargs)
  assert train_env

  eval_env = env_problem_utils.make_env(
      batch_size=eval_batch_size,
      env_problem_name=env_name,
      rendered_env=rendered_env,
      resize=resize,
      resize_dims=resize_dims,
      max_timestep=max_timestep,
      clip_rewards=clip_rewards,
      parallelism=parallelism,
      use_tpu=FLAGS.use_tpu,
      num_actions=num_actions,
      **eval_env_kwargs)
  assert eval_env

  def run_training_loop():
    """Runs the training loop."""
    logging.info('Starting the training loop.')

    trainer = trainer_class(
        output_dir=output_dir,
        train_env=train_env,
        eval_env=eval_env,
        trajectory_dump_dir=trajectory_dump_dir,
        async_mode=FLAGS.async_mode,
    )
    trainer.training_loop(n_epochs=n_epochs)

  if FLAGS.jax_debug_nans or FLAGS.disable_jit:
    with jax.disable_jit():
      run_training_loop()
  else:
    run_training_loop()

def main(argv):
  del argv
  logging.info('Starting RL training.')

  gin_configs = FLAGS.config or []
  gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)

  logging.info('Gin cofig:')
  logging.info(gin_configs)

  train_rl(
      output_dir=FLAGS.output_dir,
      train_batch_size=FLAGS.train_batch_size,
      eval_batch_size=FLAGS.eval_batch_size,
      trajectory_dump_dir=(FLAGS.trajectory_dump_dir or None),
  )

if __name__ == '__main__':
  app.run(main)
<commit_msg>RL: updated usage info<commit_after># coding=utf-8
# Copyright 2020 The Trax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

r"""Trainer for RL environments.

For now we only support PPO as RL algorithm.

Sample invocation:

TRAIN_BATCH_SIZE=32
python trax/rl_trainer.py \
  --config_file=trax/rl/configs/ppo_acrobot.gin \
  --train_batch_size=${TRAIN_BATCH_SIZE} \
  --output_dir=${HOME}/ppo_acrobot \
  --alsologtostderr
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import multiprocessing
import os

from absl import app
from absl import flags
from absl import logging
import gin
import jax
from jax.config import config
from tensor2tensor import envs  # pylint: disable=unused-import
from tensor2tensor.envs import env_problem_utils
from trax import rl  # pylint: disable=unused-import
from trax import trainer_flags  # pylint: disable=unused-import
from trax.rl import envs as rl_envs  # pylint: disable=unused-import
from trax.rl import task as rl_task
from trax.rl import trainers as rl_trainers
from trax.rl import training as light_trainers

FLAGS = flags.FLAGS

# Not just 'train' to avoid a conflict with trax.train in GIN files.
@gin.configurable(blacklist=[
    'output_dir', 'train_batch_size', 'eval_batch_size', 'trajectory_dump_dir'
])
def train_rl(
    output_dir,
    train_batch_size,
    eval_batch_size,
    env_name='Acrobot-v1',
    max_timestep=None,
    clip_rewards=False,
    rendered_env=False,
    resize=False,
    resize_dims=(105, 80),
    trainer_class=rl_trainers.PPO,
    n_epochs=10000,
    trajectory_dump_dir=None,
    num_actions=None,
    light_rl=False,
    light_rl_trainer=light_trainers.RLTrainer,
):
  """Train the RL agent.

  Args:
    output_dir: Output directory.
    train_batch_size: Number of parallel environments to use for training.
    eval_batch_size: Number of parallel environments to use for evaluation.
    env_name: Name of the environment.
    max_timestep: Int or None, the maximum number of timesteps in a trajectory.
      The environment is wrapped in a TimeLimit wrapper.
    clip_rewards: Whether to clip and discretize the rewards.
    rendered_env: Whether the environment has visual input. If so, a
      RenderedEnvProblem will be used.
    resize: whether to do resize or not
    resize_dims: Pair (height, width), dimensions to resize the visual
      observations to.
    trainer_class: RLTrainer class to use.
    n_epochs: Number epochs to run the training for.
    trajectory_dump_dir: Directory to dump trajectories to.
    num_actions: None unless one wants to use the discretization wrapper. Then
      num_actions specifies the number of discrete actions.
    light_rl: whether to use the light RL setting (experimental).
    light_rl_trainer: whichh light RL trainer to use (experimental).
  """

  if FLAGS.jax_debug_nans:
    config.update('jax_debug_nans', True)

  if FLAGS.use_tpu:
    config.update('jax_platform_name', 'tpu')
  else:
    config.update('jax_platform_name', 'gpu')

  if light_rl:
    task = rl_task.RLTask()
    trainer = light_rl_trainer(task=task, output_dir=output_dir)
    trainer.run(n_epochs)
    return

  # TODO(pkozakowski): Find a better way to determine this.
  train_env_kwargs = {}
  eval_env_kwargs = {}
  if 'OnlineTuneEnv' in env_name:
    envs_output_dir = FLAGS.envs_output_dir or os.path.join(output_dir, 'envs')
    train_env_output_dir = os.path.join(envs_output_dir, 'train')
    eval_env_output_dir = os.path.join(envs_output_dir, 'eval')
    train_env_kwargs = {'output_dir': train_env_output_dir}
    eval_env_kwargs = {'output_dir': eval_env_output_dir}

  parallelism = multiprocessing.cpu_count() if FLAGS.parallelize_envs else 1

  logging.info('Num discretized actions %s', num_actions)
  logging.info('Resize %d', resize)

  train_env = env_problem_utils.make_env(
      batch_size=train_batch_size,
      env_problem_name=env_name,
      rendered_env=rendered_env,
      resize=resize,
      resize_dims=resize_dims,
      max_timestep=max_timestep,
      clip_rewards=clip_rewards,
      parallelism=parallelism,
      use_tpu=FLAGS.use_tpu,
      num_actions=num_actions,
      **train_env_kwargs)
  assert train_env

  eval_env = env_problem_utils.make_env(
      batch_size=eval_batch_size,
      env_problem_name=env_name,
      rendered_env=rendered_env,
      resize=resize,
      resize_dims=resize_dims,
      max_timestep=max_timestep,
      clip_rewards=clip_rewards,
      parallelism=parallelism,
      use_tpu=FLAGS.use_tpu,
      num_actions=num_actions,
      **eval_env_kwargs)
  assert eval_env

  def run_training_loop():
    """Runs the training loop."""
    logging.info('Starting the training loop.')

    trainer = trainer_class(
        output_dir=output_dir,
        train_env=train_env,
        eval_env=eval_env,
        trajectory_dump_dir=trajectory_dump_dir,
        async_mode=FLAGS.async_mode,
    )
    trainer.training_loop(n_epochs=n_epochs)

  if FLAGS.jax_debug_nans or FLAGS.disable_jit:
    with jax.disable_jit():
      run_training_loop()
  else:
    run_training_loop()

def main(argv):
  del argv
  logging.info('Starting RL training.')

  gin_configs = FLAGS.config or []
  gin.parse_config_files_and_bindings(FLAGS.config_file, gin_configs)

  logging.info('Gin cofig:')
  logging.info(gin_configs)

  train_rl(
      output_dir=FLAGS.output_dir,
      train_batch_size=FLAGS.train_batch_size,
      eval_batch_size=FLAGS.eval_batch_size,
      trajectory_dump_dir=(FLAGS.trajectory_dump_dir or None),
  )

if __name__ == '__main__':
  app.run(main)
<|endoftext|>


<commit_before>#! /usr/bin/env python
#
# ===============================================================
#    Description:  Tranform SNAP graph files so that node handles
#                  are range(0, num_nodes).
#
#        Created:  12/03/2013 02:17:15 PM
#
#         Author:  Ayush Dubey, dubey@cs.cornell.edu
#
# Copyright (C) 2013, Cornell University, see the LICENSE file
#                     for licensing agreement
# ===============================================================
#

import sys
from sets import Set

def snap_to_weaver(snap_file):
    path = snap_file.split('/')
    filename = path[len(path)-1]
    path[len(path)-1] = 'weaver'
    path.append(filename)
    weaver_file = '/'.join(path)

    sfile = open(snap_file, 'r')
    nmap = {}
    edges = {}
    handle = 0
    for line in sfile:
        if len(line) == 0 or line[0] == '#':
            continue
        split = line.split()
        if len(split) != 2:
            print 'Invalid snap file'
            return
        n1 = int(split[0])
        n2 = int(split[1])
        if not n1 in nmap:
            nmap[n1] = handle
            edges[handle] = Set([])
            handle += 1
        if not n2 in nmap:
            nmap[n2] = handle
            edges[handle] = Set([])
            handle += 1
        edges[nmap[n1]].add(nmap[n2])
    sfile.close()

    wfile = open(weaver_file, 'w')
    wfile.write('#' + str(len(nmap)) + '\n')
    for n1 in edges:
        for n2 in edges[n1]:
            wfile.write(str(n1) + ' ' + str(n2) + '\n')
    wfile.close()

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print 'Needs to be called with only a single argument -- the absolute location of the snap graph file'
    else:
        snap_to_weaver(sys.argv[1])
<commit_msg>Print progress in snap cleanup script.<commit_after>#! /usr/bin/env python
#
# ===============================================================
#    Description:  Tranform SNAP graph files so that node handles
#                  are range(0, num_nodes).
#
#        Created:  12/03/2013 02:17:15 PM
#
#         Author:  Ayush Dubey, dubey@cs.cornell.edu
#
# Copyright (C) 2013, Cornell University, see the LICENSE file
#                     for licensing agreement
# ===============================================================
#

import sys
from sets import Set

def snap_to_weaver(snap_file):
    path = snap_file.split('/')
    filename = path[len(path)-1]
    path[len(path)-1] = 'weaver'
    path.append(filename)
    weaver_file = '/'.join(path)

    sfile = open(snap_file, 'r')
    nmap = {}
    edges = {}
    handle = 0
    line_count = 0
    for line in sfile:
        line_count += 1
        if line_count % 100000 == 0:
            print 'read ' + str(line_count) + ' lines'
        if len(line) == 0 or line[0] == '#':
            continue
        split = line.split()
        if len(split) != 2:
            print 'Invalid snap file'
            return
        n1 = int(split[0])
        n2 = int(split[1])
        if not n1 in nmap:
            nmap[n1] = handle
            edges[handle] = Set([])
            handle += 1
        if not n2 in nmap:
            nmap[n2] = handle
            edges[handle] = Set([])
            handle += 1
        edges[nmap[n1]].add(nmap[n2])
    sfile.close()

    wfile = open(weaver_file, 'w')
    wfile.write('#' + str(len(nmap)) + '\n')
    for n1 in edges:
        for n2 in edges[n1]:
            wfile.write(str(n1) + ' ' + str(n2) + '\n')
    wfile.close()

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print 'Needs to be called with only a single argument -- the absolute location of the snap graph file'
    else:
        snap_to_weaver(sys.argv[1])
<|endoftext|>


<commit_before>from muppet.durable_channel import DurableChannel
from muppet.remote_channel import RemoteChannel

__version__ = "0.1.4"
VERSION = tuple(map(int, __version__.split('.')))

__all__ = [
  "DurableChannel", "RemoteChannel"
]<commit_msg>Bumped version<commit_after>from muppet.durable_channel import DurableChannel
from muppet.remote_channel import RemoteChannel

__version__ = "0.1.5"
VERSION = tuple(map(int, __version__.split('.')))

__all__ = [
  "DurableChannel", "RemoteChannel"
]<|endoftext|>


<commit_before>import numpy as np
from scipy.stats import ks_2samp

from .interruptible_pool import Pool
from .clustered_kde import optimized_kde

class GetLnProbWrapper(object):
    def __init__(self, lnpost, kde):
        self.lnpost = lnpost
        self.kde = kde

    def __call__(self, p):
        lnpost = self.lnpost(p)
        kde = self.kde(p)

        return np.array([lnpost, kde])

class Sampler(object):
    """
    An Ensemble sampler.

    The :attr:`chain` member of this object has the shape:
    ``(nsteps, nwalkers, dim)`` where ``nsteps`` is the number of steps

    :param nwalkers:
        The number of individual MCMC chains to include in the ensemble.

    :param dim:
        Number of dimensions in the parameter space.

    :param lnpostfn:
        A function that takes a vector in the parameter space as input and
        returns the natural logarithm of the posterior probability for that
        position.

    """
    def __init__(self, nwalkers, ndim, lnpostfn,
                 processes=None, pool=None):
        self.nwalkers = nwalkers
        self.dim = ndim

        self._get_lnpost = lnpostfn

        self.iterations = 0

        self._kde = None

        self.pool = pool
        self.processes = processes
        if self.processes != 1 and self.pool is None:
            self.pool = Pool(self.processes)

        self._chain = np.empty((0, self.nwalkers, self.dim))
        self._lnpost = np.empty((0, self.nwalkers))
        self._lnprop = np.empty((0, self.nwalkers))
        self._acceptance = np.zeros((0, self.nwalkers))

        self._failed_p = None

    def burnin(self, p0, lnpost0=None, lnprop0=None,
               update_interval=10, max_steps=None):
        """
        Use two-sample K-S tests to determine when burnin is complete.  The
        interval over which distributions are compared will be adapted based
        on the average acceptance rate of the walkers.

        :param p0:
            A list of the initial walker positions.  It should have the
            shape ``(nwalkers, dim)``.

        :param lnpost0: (optional)
            The list of log posterior probabilities for the walkers at
            positions ``p0``. If ``lnpost0 is None``, the initial
            values are calculated. It should have the shape
            ``(nwalkers, dim)``.

        :param lnprop0: (optional)
            The list of log proposal densities for the walkers at
            positions ``p0``. If ``lnprop0 is None``, the initial
            values are calculated. It should have the shape
            ``(nwalkers, dim)``.

        :param update_interval: (optional)
            The number of steps between proposal updates.

        :param max_steps: (optional)
            An absolute maximum number of steps to take, in case burnin
            is too painful.

        After burning in, this method returns:

        * ``p`` - A list of the current walker positions, the shape of which
            will be ``(nwalkers, dim)``.

        * ``lnpost`` - The list of log posterior probabilities for the
          walkers at positions ``p``, with shape ``(nwalkers, dim)``.

        * ``lnprop`` - The list of log proposal densities for the
          walkers at positions ``p``, with shape ``(nwalkers, dim)``.

        """
        p0 = np.array(p0)

        # Go until all two-sample K-S p-values are above this
        critical_pval = 0.05

        # Start the K-S testing interval at the update interval length
        test_interval = update_interval

        # Determine the maximum iteration to look for
        start = self.iterations
        max_iter = np.inf
        if max_steps is not None:
            max_iter = start + max_steps

            # If max_steps < update interval, at least run to max_steps
            test_interval = min(test_interval, max_steps)

        burned_in = False
        while not burned_in:
            # Give up if we're about to exceed the maximum number of iterations
            if self.iterations + test_interval > max_iter:
                break

            p, lnpost, lnprop = self.sample(p0, lnpost0, lnprop0,
                                            test_interval, update_interval)

            burned_in = True
            for par in range(self.dim):
                KS, pval = ks_2samp(p0[:, par], p[:, par])

                if pval < critical_pval:
                    burned_in = False
                    break

            # Adjust the interval so >~ 90% of walkers accept a jump
            if not burned_in:
                # Use the average acceptance of the last step to window
                #   over the last 10 accepted jumps (on average)
                window = int(10 * 1/np.mean(self.acceptance[-1]))
                acceptance_rates = np.mean(self.acceptance[-window:], axis=0)

                # Use the first decile of the walkers' acceptance rates to
                #  decide the next test interval
                index = int(.1*self.nwalkers)
                low_rate = np.sort(acceptance_rates)[index]

                # If there is a lot of variance in acceptance rates, get
                #   another 10 acceptances (on average) and check again.
                if low_rate > 0.:
                    test_interval = int(1./low_rate)
                else:
                    test_interval = window

                p0, lnpost0, lnprop0 = p, lnpost, lnprop

        if not burned_in:
            print "Burnin unsuccessful."

        return (p, lnpost, lnprop)

    def sample(self, p0=None, lnpost0=None, lnprop0=None,
               iterations=1, update_interval=10):
        """
        Advance the ensemble ``iterations`` steps.

        :param p0 (optional):
            A list of the initial walker positions of shape
            ``(nwalkers, dim)``.  If ``None`` and a proposal distribution
            exists, walker positions will be drawn from the proposal.

        :param lnpost0: (optional)
            The list of log posterior probabilities for the walkers at
            positions ``p0``. If ``lnpost0 is None``, the initial
            values are calculated. It should have the shape
            ``(nwalkers, dim)``.

        :param lnprop0: (optional)
            The list of log proposal densities for the walkers at
            positions ``p0``. If ``lnprop0 is None``, the initial
            values are calculated. It should have the shape
            ``(nwalkers, dim)``.

        :param iterations: (optional)
            The number of steps to run.

        :param update_interval: (optional)
            The number of steps between proposal updates.

        After ``iteration`` steps, this method returns:

        * ``p`` - A list of the current walker positions, the shape of which
            will be ``(nwalkers, dim)``.

        * ``lnpost`` - The list of log posterior probabilities for the
          walkers at positions ``p``, with shape ``(nwalkers, dim)``.

        * ``lnprop`` - The list of log proposal densities for the
          walkers at positions ``p``, with shape ``(nwalkers, dim)``.

        """
        if p0 is None:
            p = self.draw(self.nwalkers)
        else:
            p = np.array(p0)

        if self.pool is None:
            m = map
        else:
            m = self.pool.map

        # Build a proposal if one doesn'g already exist
        if self._kde is None:
            self._kde = optimized_kde(p, pool=self.pool)

        lnpost = lnpost0
        lnprop = lnprop0

        if lnpost is None or lnprop is None:
            results = np.array(m(GetLnProbWrapper(self._get_lnpost,
                                                  self._kde), p))
            lnpost = results[:, 0] if lnpost is None else lnpost
            lnprop = results[:, 1] if lnprop is None else lnprop

        # Prepare arrays for storage ahead of time
        self._chain = np.concatenate(
            (self._chain, np.zeros((iterations, self.nwalkers, self.dim))))
        self._lnpost = np.concatenate(
            (self._lnpost, np.zeros((iterations, self.nwalkers))))
        self._lnprop = np.concatenate(
            (self._lnprop, np.zeros((iterations, self.nwalkers))))
        self._acceptance = np.concatenate(
            (self._acceptance, np.zeros((iterations, self.nwalkers))))

        for i in xrange(int(iterations)):
            try:
                # Draw new walker locations from the proposal
                p_p = self.draw(self.nwalkers)

                # Calculate the posterior probability and proposal density
                # at the proposed locations
                try:
                    results = np.array(m(GetLnProbWrapper(self._get_lnpost,
                                                          self._kde), p_p))
                    lnpost_p = results[:, 0]
                    lnprop_p = results[:, 1]

                # Catch any exceptions and exit gracefully
                except Exception as e:
                    self.rollback(self.iterations)
                    self._failed_p = p_p

                    print "Offending samples stored in ``failed_p``."
                    raise

                # Calculate the (ln) Metropolis-Hastings ration
                ln_mh_ratio = lnpost_p - lnpost + lnprop - lnprop_p

                # Accept if ratio is greater than 1
                acc = ln_mh_ratio > 0

                # Decide which of the remainder will be accepted
                worse = ~acc
                nworse = np.sum(worse)
                uniform_draws = np.random.rand(nworse)
                acc[worse] = ln_mh_ratio[worse] > np.log(uniform_draws)

                # Update locations and probability densities
                if np.any(acc):
                    p[acc] = p_p[acc]
                    lnpost[acc] = lnpost_p[acc]
                    lnprop[acc] = lnprop_p[acc]

                # Store stuff
                self._chain[self.iterations, :, :] = p
                self._lnpost[self.iterations, :] = lnpost
                self._lnprop[self.iterations, :] = lnprop
                self._acceptance[self.iterations, :] = acc

                self.iterations += 1

                # Update the proposal at the requested interval
                if self.iterations % update_interval == 0:
                    self._kde = optimized_kde(p, pool=self.pool)

            except KeyboardInterrupt:
                self.rollback(self.iterations)
                raise

        return (p, lnpost, lnprop)

    def draw(self, N):
        """
        Draw ``N`` samples from the current proposal distribution.
        """
        return self._kde.draw(N)

    @property
    def failed_p(self):
        """
        Sample that caused the last exception.
        """
        return self._failed_p

    @property
    def chain(self):
        """
        Ensemble's past samples,
            with shape ``(iterations, nwalkers, ndim)``.
        """
        return self._chain

    @property
    def lnpost(self):
        """
        Ensemble's past posterior probabilities,
            with shape ``(iterations, nwalkers)``.
        """
        return self._lnpost

    @property
    def lnprop(self):
        """
        Ensemble's past proposal probabilities,
            with shape ``(iterations, nwalkers)``.
        """
        return self._lnprop

    @property
    def acceptance(self):
        """
        Boolean array of ensemble's past acceptances,
            with shape ``(iterations, nwalkers)``.
        """
        return self._acceptance

    def rollback(self, iteration):
        """
        Shrink arrays down to a length of ``iteration`` and reset the
        pool if there is one.
        """
        self._chain = self._chain[:iteration]
        self._lnpost = self._lnpost[:iteration]
        self._lnprop = self._lnprop[:iteration]
        self._acceptance = self._acceptance[:iteration]

        # Close the old pool and open a new one
        if self.processes != 1:
            self.pool.close()
            self.pool = Pool(self.processes)
<commit_msg>Add some convenience methods for acceptance rates<commit_after>import numpy as np
from scipy.stats import ks_2samp

from .interruptible_pool import Pool
from .clustered_kde import optimized_kde

class GetLnProbWrapper(object):
    def __init__(self, lnpost, kde):
        self.lnpost = lnpost
        self.kde = kde

    def __call__(self, p):
        lnpost = self.lnpost(p)
        kde = self.kde(p)

        return np.array([lnpost, kde])

class Sampler(object):
    """
    An Ensemble sampler.

    The :attr:`chain` member of this object has the shape:
    ``(nsteps, nwalkers, dim)`` where ``nsteps`` is the number of steps

    :param nwalkers:
        The number of individual MCMC chains to include in the ensemble.

    :param dim:
        Number of dimensions in the parameter space.

    :param lnpostfn:
        A function that takes a vector in the parameter space as input and
        returns the natural logarithm of the posterior probability for that
        position.

    """
    def __init__(self, nwalkers, ndim, lnpostfn,
                 processes=None, pool=None):
        self.nwalkers = nwalkers
        self.dim = ndim

        self._get_lnpost = lnpostfn

        self.iterations = 0

        self._kde = None

        self.pool = pool
        self.processes = processes
        if self.processes != 1 and self.pool is None:
            self.pool = Pool(self.processes)

        self._chain = np.empty((0, self.nwalkers, self.dim))
        self._lnpost = np.empty((0, self.nwalkers))
        self._lnprop = np.empty((0, self.nwalkers))
        self._acceptance = np.zeros((0, self.nwalkers))

        self._failed_p = None

    def burnin(self, p0, lnpost0=None, lnprop0=None,
               update_interval=10, max_steps=None):
        """
        Use two-sample K-S tests to determine when burnin is complete.  The
        interval over which distributions are compared will be adapted based
        on the average acceptance rate of the walkers.

        :param p0:
            A list of the initial walker positions.  It should have the
            shape ``(nwalkers, dim)``.

        :param lnpost0: (optional)
            The list of log posterior probabilities for the walkers at
            positions ``p0``. If ``lnpost0 is None``, the initial
            values are calculated. It should have the shape
            ``(nwalkers, dim)``.

        :param lnprop0: (optional)
            The list of log proposal densities for the walkers at
            positions ``p0``. If ``lnprop0 is None``, the initial
            values are calculated. It should have the shape
            ``(nwalkers, dim)``.

        :param update_interval: (optional)
            The number of steps between proposal updates.

        :param max_steps: (optional)
            An absolute maximum number of steps to take, in case burnin
            is too painful.

        After burning in, this method returns:

        * ``p`` - A list of the current walker positions, the shape of which
            will be ``(nwalkers, dim)``.

        * ``lnpost`` - The list of log posterior probabilities for the
          walkers at positions ``p``, with shape ``(nwalkers, dim)``.

        * ``lnprop`` - The list of log proposal densities for the
          walkers at positions ``p``, with shape ``(nwalkers, dim)``.

        """
        p0 = np.array(p0)

        # Go until all two-sample K-S p-values are above this
        critical_pval = 0.05

        # Start the K-S testing interval at the update interval length
        test_interval = update_interval

        # Determine the maximum iteration to look for
        start = self.iterations
        max_iter = np.inf
        if max_steps is not None:
            max_iter = start + max_steps

            # If max_steps < update interval, at least run to max_steps
            test_interval = min(test_interval, max_steps)

        burned_in = False
        while not burned_in:
            # Give up if we're about to exceed the maximum number of iterations
            if self.iterations + test_interval > max_iter:
                break

            p, lnpost, lnprop = self.sample(p0, lnpost0, lnprop0,
                                            test_interval, update_interval)

            burned_in = True
            for par in range(self.dim):
                KS, pval = ks_2samp(p0[:, par], p[:, par])

                if pval < critical_pval:
                    burned_in = False
                    break

            # Adjust the interval so >~ 90% of walkers accept a jump
            if not burned_in:
                # Use the average acceptance of the last step to window
                #   over the last 10 accepted jumps (on average)
                window = int(10 * 1/np.mean(self.acceptance[-1]))
                acceptance_rates = np.mean(self.acceptance[-window:], axis=0)

                # Use the first decile of the walkers' acceptance rates to
                #  decide the next test interval
                index = int(.1*self.nwalkers)
                low_rate = np.sort(acceptance_rates)[index]

                # If there is a lot of variance in acceptance rates, get
                #   another 10 acceptances (on average) and check again.
                if low_rate > 0.:
                    test_interval = int(1./low_rate)
                else:
                    test_interval = window

                p0, lnpost0, lnprop0 = p, lnpost, lnprop

        if not burned_in:
            print "Burnin unsuccessful."

        return (p, lnpost, lnprop)

    def sample(self, p0=None, lnpost0=None, lnprop0=None,
               iterations=1, update_interval=10):
        """
        Advance the ensemble ``iterations`` steps.

        :param p0 (optional):
            A list of the initial walker positions of shape
            ``(nwalkers, dim)``.  If ``None`` and a proposal distribution
            exists, walker positions will be drawn from the proposal.

        :param lnpost0: (optional)
            The list of log posterior probabilities for the walkers at
            positions ``p0``. If ``lnpost0 is None``, the initial
            values are calculated. It should have the shape
            ``(nwalkers, dim)``.

        :param lnprop0: (optional)
            The list of log proposal densities for the walkers at
            positions ``p0``. If ``lnprop0 is None``, the initial
            values are calculated. It should have the shape
            ``(nwalkers, dim)``.

        :param iterations: (optional)
            The number of steps to run.

        :param update_interval: (optional)
            The number of steps between proposal updates.

        After ``iteration`` steps, this method returns:

        * ``p`` - A list of the current walker positions, the shape of which
            will be ``(nwalkers, dim)``.

        * ``lnpost`` - The list of log posterior probabilities for the
          walkers at positions ``p``, with shape ``(nwalkers, dim)``.

        * ``lnprop`` - The list of log proposal densities for the
          walkers at positions ``p``, with shape ``(nwalkers, dim)``.

        """
        if p0 is None:
            p = self.draw(self.nwalkers)
        else:
            p = np.array(p0)

        if self.pool is None:
            m = map
        else:
            m = self.pool.map

        # Build a proposal if one doesn'g already exist
        if self._kde is None:
            self._kde = optimized_kde(p, pool=self.pool)

        lnpost = lnpost0
        lnprop = lnprop0

        if lnpost is None or lnprop is None:
            results = np.array(m(GetLnProbWrapper(self._get_lnpost,
                                                  self._kde), p))
            lnpost = results[:, 0] if lnpost is None else lnpost
            lnprop = results[:, 1] if lnprop is None else lnprop

        # Prepare arrays for storage ahead of time
        self._chain = np.concatenate(
            (self._chain, np.zeros((iterations, self.nwalkers, self.dim))))
        self._lnpost = np.concatenate(
            (self._lnpost, np.zeros((iterations, self.nwalkers))))
        self._lnprop = np.concatenate(
            (self._lnprop, np.zeros((iterations, self.nwalkers))))
        self._acceptance = np.concatenate(
            (self._acceptance, np.zeros((iterations, self.nwalkers))))

        for i in xrange(int(iterations)):
            try:
                # Draw new walker locations from the proposal
                p_p = self.draw(self.nwalkers)

                # Calculate the posterior probability and proposal density
                # at the proposed locations
                try:
                    results = np.array(m(GetLnProbWrapper(self._get_lnpost,
                                                          self._kde), p_p))
                    lnpost_p = results[:, 0]
                    lnprop_p = results[:, 1]

                # Catch any exceptions and exit gracefully
                except Exception as e:
                    self.rollback(self.iterations)
                    self._failed_p = p_p

                    print "Offending samples stored in ``failed_p``."
                    raise

                # Calculate the (ln) Metropolis-Hastings ration
                ln_mh_ratio = lnpost_p - lnpost + lnprop - lnprop_p

                # Accept if ratio is greater than 1
                acc = ln_mh_ratio > 0

                # Decide which of the remainder will be accepted
                worse = ~acc
                nworse = np.sum(worse)
                uniform_draws = np.random.rand(nworse)
                acc[worse] = ln_mh_ratio[worse] > np.log(uniform_draws)

                # Update locations and probability densities
                if np.any(acc):
                    p[acc] = p_p[acc]
                    lnpost[acc] = lnpost_p[acc]
                    lnprop[acc] = lnprop_p[acc]

                # Store stuff
                self._chain[self.iterations, :, :] = p
                self._lnpost[self.iterations, :] = lnpost
                self._lnprop[self.iterations, :] = lnprop
                self._acceptance[self.iterations, :] = acc

                self.iterations += 1

                # Update the proposal at the requested interval
                if self.iterations % update_interval == 0:
                    self._kde = optimized_kde(p, pool=self.pool)

            except KeyboardInterrupt:
                self.rollback(self.iterations)
                raise

        return (p, lnpost, lnprop)

    def draw(self, N):
        """
        Draw ``N`` samples from the current proposal distribution.
        """
        return self._kde.draw(N)

    @property
    def failed_p(self):
        """
        Sample that caused the last exception.
        """
        return self._failed_p

    @property
    def chain(self):
        """
        Ensemble's past samples,
            with shape ``(iterations, nwalkers, ndim)``.
        """
        return self._chain

    @property
    def lnpost(self):
        """
        Ensemble's past posterior probabilities,
            with shape ``(iterations, nwalkers)``.
        """
        return self._lnpost

    @property
    def lnprop(self):
        """
        Ensemble's past proposal probabilities,
            with shape ``(iterations, nwalkers)``.
        """
        return self._lnprop

    @property
    def acceptance(self):
        """
        Boolean array of ensemble's past acceptances,
            with shape ``(iterations, nwalkers)``.
        """
        return self._acceptance

    @property
    def acceptance_fraction(self):
        """
        An array (length: ``iterations``) of the fraction of walkers that
            accepted each step.
        """
        return np.mean(self.acceptance, axis=1)

    @property
    def acceptance_rate(self, window=None):
        """
        """
        N = len(self.acceptance)

        # Use the mean acceptance rate of the last step to set the window
        if window is None:
            window = int(20 * 1.0/np.mean(self.acceptance[-1]))

        rates = np.empty((self.nwalkers, N - window + 1))
        weights = np.ones(window)/window

        for w in range(self.nwalkers):
            rates[w] = np.convolve(self.acceptance[:, w], weights, 'valid')

        return rates

    def rollback(self, iteration):
        """
        Shrink arrays down to a length of ``iteration`` and reset the
        pool if there is one.
        """
        self._chain = self._chain[:iteration]
        self._lnpost = self._lnpost[:iteration]
        self._lnprop = self._lnprop[:iteration]
        self._acceptance = self._acceptance[:iteration]

        # Close the old pool and open a new one
        if self.processes != 1:
            self.pool.close()
            self.pool = Pool(self.processes)
<|endoftext|>


<commit_before>
"""
Implementation of the command-line I{pyflakes} tool.
"""

import sys
import os
import _ast

checker = __import__('pyflakes.checker').checker

class Reporter(object):
    """
    Formats the results of pyflakes checks to users.
    """

    def __init__(self, warningStream, errorStream):
        """
        Construct a L{Reporter}.

        @param warningStream: A file-like object where warnings will be
            written to.  C{sys.stdout} is a good value.
        @param errorStream: A file-like object where error output will be
            written to.  C{sys.stderr} is a good value.
        """
        self._stdout = warningStream
        self._stderr = errorStream

    def ioError(self, filename, msg):
        """
        There was an C{IOError} while reading C{filename}.
        """
        self._stderr.write("%s: %s\n" % (filename, msg.args[1]))

    def problemDecodingSource(self, filename):
        """
        There was a problem decoding the source code in C{filename}.
        """
        self._stderr.write(filename)
        self._stderr.write(': problem decoding source\n')

    def syntaxError(self, filename, msg, lineno, offset, text):
        """
        There was a syntax errror in C{filename}.

        @param filename: The path to the file with the syntax error.
        @param msg: An explanation of the syntax error.
        @param lineno: The line number where the syntax error occurred.
        @param offset: The column on which the syntax error occurred.
        @param text: The source code containing the syntax error.
        """
        line = text.splitlines()[-1]
        if offset is not None:
            offset = offset - (len(text) - len(line))
        self._stderr.write('%s:%d: %s\n' % (filename, lineno, msg))
        self._stderr.write(line)
        self._stderr.write('\n')
        if offset is not None:
            self._stderr.write(" " * (offset + 1) + "^\n")

    def flake(self, message):
        """
        pyflakes found something wrong with the code.

        @param: A L{pyflakes.messages.Message}.
        """
        self._stdout.write(message)
        self._stdout.write('\n')

def check(codeString, filename, reporter=None):
    """
    Check the Python source given by C{codeString} for flakes.

    @param codeString: The Python source to check.
    @type codeString: C{str}

    @param filename: The name of the file the source came from, used to report
        errors.
    @type filename: C{str}

    @param reporter: A L{Reporter} instance, where errors and warnings will be
        reported.

    @return: The number of warnings emitted.
    @rtype: C{int}
    """
    if reporter is None:
        reporter = Reporter(None, sys.stderr)
    # First, compile into an AST and handle syntax errors.
    try:
        tree = compile(codeString, filename, "exec", _ast.PyCF_ONLY_AST)
    except SyntaxError, value:
        msg = value.args[0]

        (lineno, offset, text) = value.lineno, value.offset, value.text

        # If there's an encoding problem with the file, the text is None.
        if text is None:
            # Avoid using msg, since for the only known case, it contains a
            # bogus message that claims the encoding the file declared was
            # unknown.
            reporter.problemDecodingSource(filename)
        else:
            reporter.syntaxError(filename, msg, lineno, offset, text)
        return 1
    else:
        # Okay, it's syntactically valid.  Now check it.
        w = checker.Checker(tree, filename)
        w.messages.sort(lambda a, b: cmp(a.lineno, b.lineno))
        for warning in w.messages:
            reporter.flake(warning)
        return len(w.messages)

def checkPath(filename, reporter=None):
    """
    Check the given path, printing out any warnings detected.

    @param reporter: A L{Reporter} instance, where errors and warnings will be
        reported.

    @return: the number of warnings printed
    """
    if reporter is None:
        reporter = Reporter(None, sys.stderr)
    try:
        return check(file(filename, 'U').read() + '\n', filename, reporter)
    except IOError, msg:
        reporter.ioError(filename, msg)
        return 1

def iterSourceCode(paths):
    """
    Iterate over all Python source files in C{paths}.

    @param paths: A list of paths.  Directories will be recursed into and
        any .py files found will be yielded.  Any non-directories will be
        yielded as-is.
    """
    for path in paths:
        if os.path.isdir(path):
            for dirpath, dirnames, filenames in os.walk(path):
                for filename in filenames:
                    if filename.endswith('.py'):
                        yield os.path.join(dirpath, filename)
        else:
            yield path

def main():
    warnings = 0
    args = sys.argv[1:]
    reporter = Reporter(None, sys.stderr)
    if args:
        for arg in args:
            if os.path.isdir(arg):
                for dirpath, dirnames, filenames in os.walk(arg):
                    for filename in filenames:
                        if filename.endswith('.py'):
                            warnings += checkPath(
                                os.path.join(dirpath, filename), reporter)
            else:
                warnings += checkPath(arg, reporter)
    else:
        warnings += check(sys.stdin.read(), '<stdin>')

    raise SystemExit(warnings > 0)
<commit_msg>Now we have iterSourceCode instead.<commit_after>
"""
Implementation of the command-line I{pyflakes} tool.
"""

import sys
import os
import _ast

checker = __import__('pyflakes.checker').checker

class Reporter(object):
    """
    Formats the results of pyflakes checks to users.
    """

    def __init__(self, warningStream, errorStream):
        """
        Construct a L{Reporter}.

        @param warningStream: A file-like object where warnings will be
            written to.  C{sys.stdout} is a good value.
        @param errorStream: A file-like object where error output will be
            written to.  C{sys.stderr} is a good value.
        """
        self._stdout = warningStream
        self._stderr = errorStream

    def ioError(self, filename, msg):
        """
        There was an C{IOError} while reading C{filename}.
        """
        self._stderr.write("%s: %s\n" % (filename, msg.args[1]))

    def problemDecodingSource(self, filename):
        """
        There was a problem decoding the source code in C{filename}.
        """
        self._stderr.write(filename)
        self._stderr.write(': problem decoding source\n')

    def syntaxError(self, filename, msg, lineno, offset, text):
        """
        There was a syntax errror in C{filename}.

        @param filename: The path to the file with the syntax error.
        @param msg: An explanation of the syntax error.
        @param lineno: The line number where the syntax error occurred.
        @param offset: The column on which the syntax error occurred.
        @param text: The source code containing the syntax error.
        """
        line = text.splitlines()[-1]
        if offset is not None:
            offset = offset - (len(text) - len(line))
        self._stderr.write('%s:%d: %s\n' % (filename, lineno, msg))
        self._stderr.write(line)
        self._stderr.write('\n')
        if offset is not None:
            self._stderr.write(" " * (offset + 1) + "^\n")

    def flake(self, message):
        """
        pyflakes found something wrong with the code.

        @param: A L{pyflakes.messages.Message}.
        """
        self._stdout.write(message)
        self._stdout.write('\n')

def check(codeString, filename, reporter=None):
    """
    Check the Python source given by C{codeString} for flakes.

    @param codeString: The Python source to check.
    @type codeString: C{str}

    @param filename: The name of the file the source came from, used to report
        errors.
    @type filename: C{str}

    @param reporter: A L{Reporter} instance, where errors and warnings will be
        reported.

    @return: The number of warnings emitted.
    @rtype: C{int}
    """
    if reporter is None:
        reporter = Reporter(None, sys.stderr)
    # First, compile into an AST and handle syntax errors.
    try:
        tree = compile(codeString, filename, "exec", _ast.PyCF_ONLY_AST)
    except SyntaxError, value:
        msg = value.args[0]

        (lineno, offset, text) = value.lineno, value.offset, value.text

        # If there's an encoding problem with the file, the text is None.
        if text is None:
            # Avoid using msg, since for the only known case, it contains a
            # bogus message that claims the encoding the file declared was
            # unknown.
            reporter.problemDecodingSource(filename)
        else:
            reporter.syntaxError(filename, msg, lineno, offset, text)
        return 1
    else:
        # Okay, it's syntactically valid.  Now check it.
        w = checker.Checker(tree, filename)
        w.messages.sort(lambda a, b: cmp(a.lineno, b.lineno))
        for warning in w.messages:
            reporter.flake(warning)
        return len(w.messages)

def checkPath(filename, reporter=None):
    """
    Check the given path, printing out any warnings detected.

    @param reporter: A L{Reporter} instance, where errors and warnings will be
        reported.

    @return: the number of warnings printed
    """
    if reporter is None:
        reporter = Reporter(None, sys.stderr)
    try:
        return check(file(filename, 'U').read() + '\n', filename, reporter)
    except IOError, msg:
        reporter.ioError(filename, msg)
        return 1

def iterSourceCode(paths):
    """
    Iterate over all Python source files in C{paths}.

    @param paths: A list of paths.  Directories will be recursed into and
        any .py files found will be yielded.  Any non-directories will be
        yielded as-is.
    """
    for path in paths:
        if os.path.isdir(path):
            for dirpath, dirnames, filenames in os.walk(path):
                for filename in filenames:
                    if filename.endswith('.py'):
                        yield os.path.join(dirpath, filename)
        else:
            yield path

def main():
    warnings = 0
    args = sys.argv[1:]
    reporter = Reporter(None, sys.stderr)
    if args:
        for sourcePath in iterSourceCode(args):
            warnings += checkPath(sourcePath, reporter)
    else:
        warnings += check(sys.stdin.read(), '<stdin>')

    raise SystemExit(warnings > 0)
<|endoftext|>