<commit_before>import os
import os.path
import inspect
import numpy as np
from astropy.io import fits
import yaml
import galsim
import bashes

class Observation(object):
    """
    Represents a GREAT3 observation specified by a branch, index (0-199) and epoch.
    """
    def __init__(self,branch,index,epoch):
        """
        Initializes an observation using a branch path of the form 'control/ground/constant'
        that should be present under $GREAT3_ROOT (and also under $GREAT3_ROOT/truth
        if truth info is required), together with an image index (0-199) and an epoch number.
        Raises a RuntimeError if any problems are detected. After initialization, the following
        attributes are defined: nFields, nSubfieldsPerField, nEpochs, pixelScale, stampSize.
        """
        # Lookup the GREAT3 filesystem root.
        if 'GREAT3_ROOT' not in os.environ:
            raise RuntimeError('$GREAT3_ROOT is not set.')
        g3root = os.environ['GREAT3_ROOT']
        # Check for a valid branch path.
        pathNames = branch.split('/')
        if (len(pathNames) != 3 or
            pathNames[0] not in ('control','real_galaxy','variable_psf','multiepoch','full') or
            pathNames[1] not in ('ground','space') or
            pathNames[2] not in ('constant','variable')):
            raise RuntimeError('Invalid branch path: %r' % branch)
        # Lookup the path to this observation's branch.
        self.branchPath = os.path.join(g3root,branch)
        if not os.path.isdir(self.branchPath):
            raise RuntimeError('No such branch path: %r' % self.branchPath)
        # Do we have truth info available?
        self.truthPath = os.path.join(g3root,'truth',branch)
        if not os.path.isdir(self.truthPath):
            self.truthPath = None
        # Specify this branch's parameters.
        if pathNames[0] in ('variable_psf','full') or pathNames[2] == 'variable':
            self.nFields = 10
            self.nSubfieldsPerField = 20
        else:
            self.nFields = 200
            self.nSubfieldsPerField = 1
        if pathNames[0] in ('multiepoch','full'):
            self.nEpochs = 6
        else:
            self.nEpochs = 1
        if pathNames[1] == 'space':
            if self.nEpochs == 1:
                self.pixelScale = 0.05
                self.stampSize = 96
            else:
                self.pixelScale = 0.1
                self.stampSize = 48
        else:
            self.pixelScale = 0.2
            self.stampSize = 48
        # Check the index and epoch parameters.
        try:
            self.index = int(index)
            assert self.index >= 0 and self.index < 200
        except (ValueError,AssertionError):
            raise RuntimeError('Invalid branch index: %r' % index)
        try:
            self.epoch = int(epoch)
            assert self.epoch >= 0 and self.epoch < self.nEpochs
        except (ValueError,AssertionError):
            raise RuntimeError('Invalid branch epoch index: %r' % epoch)
        # Our galaxy and star images are loaded on demand.
        self.image = None
        self.stars = None
        # Our truth params and catalog are loaded on demand.
        self.truthParams = None
        self.truthCatalog = None

    @staticmethod
    def addArgs(parser):
        """
        Add arguments to the provided command-line parser that support the fromArgs() method.
        """
        parser.add_argument('--branch', type = str, default = 'control/ground/constant',
            help = 'Name of branch to use relative to $GREAT3_ROOT')
        parser.add_argument('--index', type = int, default = 0,
            help = 'Index of field to analyze (0-199)')
        parser.add_argument('--epoch', type = int, default = 0,
            help = 'Epoch number to analyze')

    @staticmethod
    def fromArgs(args):
        """
        Returns a dictionary of constructor parameter values based on the parsed args provided.
        """
        # Look up the named Estimator constructor parameters.
        pnames = (inspect.getargspec(Observation.__init__)).args[1:]
        # Get a dictionary of the arguments provided.
        argsDict = vars(args)
        # Return a dictionary of constructor parameters provided in args.
        return { key:argsDict[key] for key in (set(pnames) & set(argsDict)) }

    @classmethod
    def getGSParams(cls):
        if not hasattr(cls,'GSParams'):
            cls.GSParams = galsim.GSParams(maximum_fft_size=2**16)
        return cls.GSParams

    def getImage(self):
        """
        Returns the array of postage stamp image data for this observation and initializes
        our stampSize data member.
        """
        if self.image is None:
            dataStampsPath = os.path.join(self.branchPath,'image-%03d-%d.fits' % (
                self.index,self.epoch))
            hduList = fits.open(dataStampsPath)
            dataStamps = hduList[0].data
            hduList.close()
            # Check for the expected image dimensions.
            assert dataStamps.shape[0] == dataStamps.shape[1], 'Image data is not square'
            assert dataStamps.shape[0] == 100*self.stampSize, 'Image has unexpected dimensions'
            self.image = bashes.tiled.Tiled(dataStamps,self.stampSize)
            self.image.scale = self.pixelScale
        return self.image

    def getStars(self):
        """
        Returns the array of postage stamp starfield data for this observation.
        """
        if self.stars is None:
            psfStampsPath = os.path.join(self.branchPath,'starfield_image-%03d-%d.fits' % (
                self.index,self.epoch))
            hduList = fits.open(psfStampsPath)
            psfStamps = hduList[0].data
            hduList.close()
            self.stars = bashes.tiled.Tiled(psfStamps,self.stampSize)
            self.stars.scale = self.pixelScale
        return self.stars

    def getTruthParams(self):
        """
        Returns a dictionary of truth parameter values for this observation.
        """
        if self.truthParams is None:
            # No cached value available, so fetch it now.
            if not self.truthPath:
                raise RuntimeError('No truth available for observation')
            truthParamsPath = os.path.join(self.truthPath,'epoch_parameters-%03d-%d.yaml' % (
                self.index,self.epoch))
            with open(truthParamsPath,'r') as f:
                self.truthParams = yaml.load(f)
        return self.truthParams

    def getTruthCatalog(self):
        """
        Returns the truth catalog for this observation.
        """
        if self.truthCatalog is None:
            # No cached value available, so fetch it now.
            if not self.truthPath:
                raise RuntimeError('No truth available for observation')
            truthCatalogPath = os.path.join(self.truthPath,'epoch_catalog-%03d-%d.fits' % (
                self.index,self.epoch))
            hduList = fits.open(truthCatalogPath)
            self.truthCatalog = hduList[1].data
            hduList.close()
        return self.truthCatalog

    def createSource(self,galaxyIndex,shifted = False,lensed = False):
        """
        Returns a GalSim model of the source for the specified galaxy index with
        optional centroid shifts and weak lensing distortion.
        """
        params = self.getTruthCatalog()[galaxyIndex]
        # Create the bulge component.
        bulge = galsim.Sersic(flux = params['bulge_flux'],
            half_light_radius = params['bulge_hlr'],
            n = params['bulge_n'], gsparams = Observation.getGSParams())
        bulge.applyShear(q = params['bulge_q'],
            beta = params['bulge_beta_radians']*galsim.radians)
        # Is there a disk component?
        if params['disk_flux'] > 0:
            disk = galsim.Exponential(flux = params['disk_flux'],
                half_light_radius = params['disk_hlr'],
                gsparams = Observation.getGSParams())
            disk.applyShear(q = params['disk_q'],
                beta = params['disk_beta_radians']*galsim.radians)
            source = galsim.Add(bulge,disk)
        else:
            source = bulge
        # Apply optional centroid shift.
        if shifted:
            source = source.shift(
                dx=params['xshift']*self.pixelScale,
                dy=params['yshift']*self.pixelScale)
        # Apply optional lensing.
        if lensed:
            source = source.lens(g1=params['g1'],g2=params['g2'],mu=params['mu'])
        return source

    def createPSF(self,galaxyIndex):
        """
        Returns a GalSim model of the PSF for the specified galaxy index.
        """
        catalog = self.getTruthCatalog()
        keys = catalog.columns.names
        params = catalog[galaxyIndex]
        # Create an empty list of models that will be convolved for the final PSF.
        models = [ ]
        # Add jitter contribution if provided.
        if 'opt_psf_jitter_sigma' in keys:
            jitterPSF = galsim.Gaussian(sigma=params['opt_psf_jitter_sigma']).shear(
                beta = params['opt_psf_jitter_beta']*galsim.degrees,
                e = params['opt_psf_jitter_e'])
            models.append(jitterPSF)
        # Add charge diffusion contribution if provided.
        if 'opt_psf_charge_sigma' in keys:
            chargePSF = galsim.Gaussian(sigma=params['opt_psf_charge_sigma']).shear(
                e1 = params['opt_psf_charge_e1'], e2 = 0.)
            models.append(chargePSF)
        # Create the optical component, which is always present.
        kmap = { 'opt_psf_lam_over_diam':'lam_over_diam', 'opt_psf_obscuration':'obscuration',
            'opt_psf_n_struts':'nstruts', 'opt_psf_strut_angle':'strut_angle',
            'opt_psf_pad_factor':'pad_factor', 'opt_psf_defocus':'defocus',
            'opt_psf_astig1':'astig1', 'opt_psf_astig2':'astig2', 'opt_psf_coma1':'coma1',
            'opt_psf_coma2':'coma2', 'opt_psf_trefoil1':'trefoil1', 'opt_psf_trefoil2':'trefoil2',
            'opt_psf_spher':'spher'}
        opticalPSFParams = { kmap[key]:params[key] for key in kmap }
        # Add units for the strut angle.
        opticalPSFParams['strut_angle'] *= galsim.degrees
        # Suppress warnings.
        opticalPSFParams['suppress_warning'] = True
        # Build the optical PSF from the params dictionary.
        models.append(galsim.OpticalPSF(**opticalPSFParams))
        # Add an atmospheric component if a FWHM value is provided.
        if 'atmos_psf_fwhm' in keys:
            atmosphericPSF = galsim.Kolmogorov(fwhm = params['atmos_psf_fwhm']).shear(
                beta = params['atmos_psf_beta']*galsim.degrees,
                e = params['atmos_psf_e'])
            models.append(atmosphericPSF)
        # Return the convolution of all PSF component models.
        return galsim.Convolve(models, gsparams = Observation.getGSParams())

    def createObject(self,galaxyIndex,shifted = True,lensed = True):
        """
        Returns a GalSim model of the object corresponding to the specified galaxy index,
        consisting of the source model with lensing distortion and centroid shift applied,
        and convolved with the appropriate PSF.
        """
        # Look up the component models.
        src = self.createSource(galaxyIndex,shifted,lensed)
        psf = self.createPSF(galaxyIndex)
        # Return their convolution.
        return galsim.Convolve(src,psf)

    def renderObject(self,galaxyIndex,shifted = True,lensed = True,addNoise = True):
        """
        Renders a postage stamp of the truth model for the specified galaxy index
        with optional noise (that will exactly match the noise used for GREAT3).
        """
        obj = self.createObject(galaxyIndex)
        stamp = bashes.utility.render(obj,self.pixelScale,size = self.stampSize)
        if addNoise:
            params = self.getTruthParams()
            seed = params['noise_seed']
            var = float(params['noise']['variance'])
            rng = galsim.BaseDeviate(seed = seed + galaxyIndex)
            noise = galsim.GaussianNoise(rng).withVariance(var)
            stamp.addNoise(noise)
        return stamp

def main():

    import argparse
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    Observation.addArgs(parser)
    parser.add_argument('--stamp', type = int, default = 0,
        help = 'index of stamp to use (0-99999)')
    parser.add_argument('--unlensed', action = 'store_true',
        help = 'do not include weak lensing effects')
    parser.add_argument('--test', action = 'store_true',
        help = 'test that stamp reconstructed from truth matches actual stamp')
    parser.add_argument('--ds9', action = 'store_true',
        help = 'display stamps in DS9')
    parser.add_argument('--save', type = str, default = None,
        help = 'save stamps to the specified FITS file')
    parser.add_argument('--truth', action = 'store_true',
        help = 'print catalog truth parameter values')
    args = parser.parse_args()

    # Initialize the requested observation.
    obs = Observation(**Observation.fromArgs(args))

    # Dump catalog truth info for this stamp if requested.
    if args.truth:
        import pprint
        catalog = obs.getTruthCatalog()
        truth = zip(catalog.columns.names,catalog[args.stamp])
        pprint.pprint(truth)

    # Render stamps if requested.
    if args.ds9 or args.save or args.test:
        lensed = not args.unlensed
        # Lookup the specified stamp's psf and source models.
        psfModel = obs.createPSF(args.stamp)
        srcModel = obs.createSource(args.stamp,shifted=True,lensed=lensed)

        # Render the PSF and source models separately.
        gsp = galsim.GSParams(maximum_fft_size = 2**16)
        psfStamp = bashes.utility.render(psfModel,obs.pixelScale,size=obs.stampSize)
        srcStamp = bashes.utility.render(srcModel,obs.pixelScale,size=obs.stampSize)

        # Render the combined object with and without noise.
        objStamp = obs.renderObject(args.stamp,shifted=True,lensed=lensed,addNoise=False)
        noiseStamp = obs.renderObject(args.stamp,shifted=True,lensed=lensed,addNoise=True)

        if args.test:
            dataStamp = obs.getImage().getStamp(args.stamp)

        if args.ds9:
            display = bashes.display.Display('cmap heat; scale sqrt')
            display.show(psfStamp)
            display.show(srcStamp)
            display.show(noiseStamp)
            if args.test:
                display.show(dataStamp,reuseLimits=True)
            display.show(objStamp,reuseLimits=True)

        if args.save:
            # Open this file using: ds9 -zoom to fit -color heat -multiframe <filename>
            stamps = [objStamp,psfStamp,srcStamp,noiseStamp]
            if args.test:
                stamps.append(dataStamp)
            galsim.fits.writeMulti(stamps, file_name = args.save)

        if args.test:
            delta = noiseStamp.array - dataStamp.array
            adiff = np.max(np.abs(delta))
            nonzero = dataStamp.array != 0
            rdiff = np.max(np.abs(delta[nonzero]/dataStamp.array[nonzero]))
            print 'Max difference between generated and saved stamps: %.3g (abs) %.3g (rel)' % (
                adiff,rdiff)
            noiseVar = obs.getTruthParams()['noise']['variance']
            print 'Std. deviation of differences / noise RMS = %.3g' % (
                np.std(delta)/np.sqrt(noiseVar))

if __name__ == "__main__":
    main()

<commit_msg>Swap order of lensing and centroid shift since lens operation is relative to stamp center, not profile centroid<commit_after>import os
import os.path
import inspect
import numpy as np
from astropy.io import fits
import yaml
import galsim
import bashes

class Observation(object):
    """
    Represents a GREAT3 observation specified by a branch, index (0-199) and epoch.
    """
    def __init__(self,branch,index,epoch):
        """
        Initializes an observation using a branch path of the form 'control/ground/constant'
        that should be present under $GREAT3_ROOT (and also under $GREAT3_ROOT/truth
        if truth info is required), together with an image index (0-199) and an epoch number.
        Raises a RuntimeError if any problems are detected. After initialization, the following
        attributes are defined: nFields, nSubfieldsPerField, nEpochs, pixelScale, stampSize.
        """
        # Lookup the GREAT3 filesystem root.
        if 'GREAT3_ROOT' not in os.environ:
            raise RuntimeError('$GREAT3_ROOT is not set.')
        g3root = os.environ['GREAT3_ROOT']
        # Check for a valid branch path.
        pathNames = branch.split('/')
        if (len(pathNames) != 3 or
            pathNames[0] not in ('control','real_galaxy','variable_psf','multiepoch','full') or
            pathNames[1] not in ('ground','space') or
            pathNames[2] not in ('constant','variable')):
            raise RuntimeError('Invalid branch path: %r' % branch)
        # Lookup the path to this observation's branch.
        self.branchPath = os.path.join(g3root,branch)
        if not os.path.isdir(self.branchPath):
            raise RuntimeError('No such branch path: %r' % self.branchPath)
        # Do we have truth info available?
        self.truthPath = os.path.join(g3root,'truth',branch)
        if not os.path.isdir(self.truthPath):
            self.truthPath = None
        # Specify this branch's parameters.
        if pathNames[0] in ('variable_psf','full') or pathNames[2] == 'variable':
            self.nFields = 10
            self.nSubfieldsPerField = 20
        else:
            self.nFields = 200
            self.nSubfieldsPerField = 1
        if pathNames[0] in ('multiepoch','full'):
            self.nEpochs = 6
        else:
            self.nEpochs = 1
        if pathNames[1] == 'space':
            if self.nEpochs == 1:
                self.pixelScale = 0.05
                self.stampSize = 96
            else:
                self.pixelScale = 0.1
                self.stampSize = 48
        else:
            self.pixelScale = 0.2
            self.stampSize = 48
        # Check the index and epoch parameters.
        try:
            self.index = int(index)
            assert self.index >= 0 and self.index < 200
        except (ValueError,AssertionError):
            raise RuntimeError('Invalid branch index: %r' % index)
        try:
            self.epoch = int(epoch)
            assert self.epoch >= 0 and self.epoch < self.nEpochs
        except (ValueError,AssertionError):
            raise RuntimeError('Invalid branch epoch index: %r' % epoch)
        # Our galaxy and star images are loaded on demand.
        self.image = None
        self.stars = None
        # Our truth params and catalog are loaded on demand.
        self.truthParams = None
        self.truthCatalog = None

    @staticmethod
    def addArgs(parser):
        """
        Add arguments to the provided command-line parser that support the fromArgs() method.
        """
        parser.add_argument('--branch', type = str, default = 'control/ground/constant',
            help = 'Name of branch to use relative to $GREAT3_ROOT')
        parser.add_argument('--index', type = int, default = 0,
            help = 'Index of field to analyze (0-199)')
        parser.add_argument('--epoch', type = int, default = 0,
            help = 'Epoch number to analyze')

    @staticmethod
    def fromArgs(args):
        """
        Returns a dictionary of constructor parameter values based on the parsed args provided.
        """
        # Look up the named Estimator constructor parameters.
        pnames = (inspect.getargspec(Observation.__init__)).args[1:]
        # Get a dictionary of the arguments provided.
        argsDict = vars(args)
        # Return a dictionary of constructor parameters provided in args.
        return { key:argsDict[key] for key in (set(pnames) & set(argsDict)) }

    @classmethod
    def getGSParams(cls):
        if not hasattr(cls,'GSParams'):
            cls.GSParams = galsim.GSParams(maximum_fft_size=2**16)
        return cls.GSParams

    def getImage(self):
        """
        Returns the array of postage stamp image data for this observation and initializes
        our stampSize data member.
        """
        if self.image is None:
            dataStampsPath = os.path.join(self.branchPath,'image-%03d-%d.fits' % (
                self.index,self.epoch))
            hduList = fits.open(dataStampsPath)
            dataStamps = hduList[0].data
            hduList.close()
            # Check for the expected image dimensions.
            assert dataStamps.shape[0] == dataStamps.shape[1], 'Image data is not square'
            assert dataStamps.shape[0] == 100*self.stampSize, 'Image has unexpected dimensions'
            self.image = bashes.tiled.Tiled(dataStamps,self.stampSize)
            self.image.scale = self.pixelScale
        return self.image

    def getStars(self):
        """
        Returns the array of postage stamp starfield data for this observation.
        """
        if self.stars is None:
            psfStampsPath = os.path.join(self.branchPath,'starfield_image-%03d-%d.fits' % (
                self.index,self.epoch))
            hduList = fits.open(psfStampsPath)
            psfStamps = hduList[0].data
            hduList.close()
            self.stars = bashes.tiled.Tiled(psfStamps,self.stampSize)
            self.stars.scale = self.pixelScale
        return self.stars

    def getTruthParams(self):
        """
        Returns a dictionary of truth parameter values for this observation.
        """
        if self.truthParams is None:
            # No cached value available, so fetch it now.
            if not self.truthPath:
                raise RuntimeError('No truth available for observation')
            truthParamsPath = os.path.join(self.truthPath,'epoch_parameters-%03d-%d.yaml' % (
                self.index,self.epoch))
            with open(truthParamsPath,'r') as f:
                self.truthParams = yaml.load(f)
        return self.truthParams

    def getTruthCatalog(self):
        """
        Returns the truth catalog for this observation.
        """
        if self.truthCatalog is None:
            # No cached value available, so fetch it now.
            if not self.truthPath:
                raise RuntimeError('No truth available for observation')
            truthCatalogPath = os.path.join(self.truthPath,'epoch_catalog-%03d-%d.fits' % (
                self.index,self.epoch))
            hduList = fits.open(truthCatalogPath)
            self.truthCatalog = hduList[1].data
            hduList.close()
        return self.truthCatalog

    def createSource(self,galaxyIndex,shifted = False,lensed = False):
        """
        Returns a GalSim model of the source for the specified galaxy index with
        optional centroid shifts and weak lensing distortion.
        """
        params = self.getTruthCatalog()[galaxyIndex]
        # Create the bulge component.
        bulge = galsim.Sersic(flux = params['bulge_flux'],
            half_light_radius = params['bulge_hlr'],
            n = params['bulge_n'], gsparams = Observation.getGSParams())
        bulge.applyShear(q = params['bulge_q'],
            beta = params['bulge_beta_radians']*galsim.radians)
        # Is there a disk component?
        if params['disk_flux'] > 0:
            disk = galsim.Exponential(flux = params['disk_flux'],
                half_light_radius = params['disk_hlr'],
                gsparams = Observation.getGSParams())
            disk.applyShear(q = params['disk_q'],
                beta = params['disk_beta_radians']*galsim.radians)
            source = galsim.Add(bulge,disk)
        else:
            source = bulge
        # Apply optional lensing.
        if lensed:
            source = source.lens(g1=params['g1'],g2=params['g2'],mu=params['mu'])
        # Apply optional centroid shift.
        if shifted:
            source = source.shift(
                dx=params['xshift']*self.pixelScale,
                dy=params['yshift']*self.pixelScale)
        return source

    def createPSF(self,galaxyIndex):
        """
        Returns a GalSim model of the PSF for the specified galaxy index.
        """
        catalog = self.getTruthCatalog()
        keys = catalog.columns.names
        params = catalog[galaxyIndex]
        # Create an empty list of models that will be convolved for the final PSF.
        models = [ ]
        # Add jitter contribution if provided.
        if 'opt_psf_jitter_sigma' in keys:
            jitterPSF = galsim.Gaussian(sigma=params['opt_psf_jitter_sigma']).shear(
                beta = params['opt_psf_jitter_beta']*galsim.degrees,
                e = params['opt_psf_jitter_e'])
            models.append(jitterPSF)
        # Add charge diffusion contribution if provided.
        if 'opt_psf_charge_sigma' in keys:
            chargePSF = galsim.Gaussian(sigma=params['opt_psf_charge_sigma']).shear(
                e1 = params['opt_psf_charge_e1'], e2 = 0.)
            models.append(chargePSF)
        # Create the optical component, which is always present.
        kmap = { 'opt_psf_lam_over_diam':'lam_over_diam', 'opt_psf_obscuration':'obscuration',
            'opt_psf_n_struts':'nstruts', 'opt_psf_strut_angle':'strut_angle',
            'opt_psf_pad_factor':'pad_factor', 'opt_psf_defocus':'defocus',
            'opt_psf_astig1':'astig1', 'opt_psf_astig2':'astig2', 'opt_psf_coma1':'coma1',
            'opt_psf_coma2':'coma2', 'opt_psf_trefoil1':'trefoil1', 'opt_psf_trefoil2':'trefoil2',
            'opt_psf_spher':'spher'}
        opticalPSFParams = { kmap[key]:params[key] for key in kmap }
        # Add units for the strut angle.
        opticalPSFParams['strut_angle'] *= galsim.degrees
        # Suppress warnings.
        opticalPSFParams['suppress_warning'] = True
        # Build the optical PSF from the params dictionary.
        models.append(galsim.OpticalPSF(**opticalPSFParams))
        # Add an atmospheric component if a FWHM value is provided.
        if 'atmos_psf_fwhm' in keys:
            atmosphericPSF = galsim.Kolmogorov(fwhm = params['atmos_psf_fwhm']).shear(
                beta = params['atmos_psf_beta']*galsim.degrees,
                e = params['atmos_psf_e'])
            models.append(atmosphericPSF)
        # Return the convolution of all PSF component models.
        return galsim.Convolve(models, gsparams = Observation.getGSParams())

    def createObject(self,galaxyIndex,shifted = True,lensed = True):
        """
        Returns a GalSim model of the object corresponding to the specified galaxy index,
        consisting of the source model with lensing distortion and centroid shift applied,
        and convolved with the appropriate PSF.
        """
        # Look up the component models.
        src = self.createSource(galaxyIndex,shifted,lensed)
        psf = self.createPSF(galaxyIndex)
        # Return their convolution.
        return galsim.Convolve(src,psf)

    def renderObject(self,galaxyIndex,shifted = True,lensed = True,addNoise = True):
        """
        Renders a postage stamp of the truth model for the specified galaxy index
        with optional noise (that will exactly match the noise used for GREAT3).
        """
        obj = self.createObject(galaxyIndex)
        stamp = bashes.utility.render(obj,self.pixelScale,size = self.stampSize)
        if addNoise:
            params = self.getTruthParams()
            seed = params['noise_seed']
            var = float(params['noise']['variance'])
            rng = galsim.BaseDeviate(seed = seed + galaxyIndex)
            noise = galsim.GaussianNoise(rng).withVariance(var)
            stamp.addNoise(noise)
        return stamp

def main():

    import argparse
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    Observation.addArgs(parser)
    parser.add_argument('--stamp', type = int, default = 0,
        help = 'index of stamp to use (0-99999)')
    parser.add_argument('--unlensed', action = 'store_true',
        help = 'do not include weak lensing effects')
    parser.add_argument('--test', action = 'store_true',
        help = 'test that stamp reconstructed from truth matches actual stamp')
    parser.add_argument('--ds9', action = 'store_true',
        help = 'display stamps in DS9')
    parser.add_argument('--save', type = str, default = None,
        help = 'save stamps to the specified FITS file')
    parser.add_argument('--truth', action = 'store_true',
        help = 'print catalog truth parameter values')
    args = parser.parse_args()

    # Initialize the requested observation.
    obs = Observation(**Observation.fromArgs(args))

    # Dump catalog truth info for this stamp if requested.
    if args.truth:
        import pprint
        catalog = obs.getTruthCatalog()
        truth = zip(catalog.columns.names,catalog[args.stamp])
        pprint.pprint(truth)

    # Render stamps if requested.
    if args.ds9 or args.save or args.test:
        lensed = not args.unlensed
        # Lookup the specified stamp's psf and source models.
        psfModel = obs.createPSF(args.stamp)
        srcModel = obs.createSource(args.stamp,shifted=True,lensed=lensed)

        # Render the PSF and source models separately.
        gsp = galsim.GSParams(maximum_fft_size = 2**16)
        psfStamp = bashes.utility.render(psfModel,obs.pixelScale,size=obs.stampSize)
        srcStamp = bashes.utility.render(srcModel,obs.pixelScale,size=obs.stampSize)

        # Render the combined object with and without noise.
        objStamp = obs.renderObject(args.stamp,shifted=True,lensed=lensed,addNoise=False)
        noiseStamp = obs.renderObject(args.stamp,shifted=True,lensed=lensed,addNoise=True)

        if args.test:
            dataStamp = obs.getImage().getStamp(args.stamp)

        if args.ds9:
            display = bashes.display.Display('cmap heat; scale sqrt')
            display.show(psfStamp)
            display.show(srcStamp)
            display.show(noiseStamp)
            if args.test:
                display.show(dataStamp,reuseLimits=True)
            display.show(objStamp,reuseLimits=True)

        if args.save:
            # Open this file using: ds9 -zoom to fit -color heat -multiframe <filename>
            stamps = [objStamp,psfStamp,srcStamp,noiseStamp]
            if args.test:
                stamps.append(dataStamp)
            galsim.fits.writeMulti(stamps, file_name = args.save)

        if args.test:
            delta = noiseStamp.array - dataStamp.array
            adiff = np.max(np.abs(delta))
            nonzero = dataStamp.array != 0
            rdiff = np.max(np.abs(delta[nonzero]/dataStamp.array[nonzero]))
            print 'Max difference between generated and saved stamps: %.3g (abs) %.3g (rel)' % (
                adiff,rdiff)
            noiseVar = obs.getTruthParams()['noise']['variance']
            print 'Std. deviation of differences / noise RMS = %.3g' % (
                np.std(delta)/np.sqrt(noiseVar))

if __name__ == "__main__":
    main()

<|endoftext|>


<commit_before>import time
import os
from constants import constants
import subprocess
from constants import blockNames
import shutil
import simulationTypes
import string
import sys

class SimulationParameters:
    def __init__(self):
        self.processors = 0
        self.executable = ""
        self.templetizedControlFile = ""
        self.dataFolder = ""
        self.documentsFolder = ""
        self.iterations = 0
        self.peleSteps = 0
        self.seed = 0
        self.exitCondition = None

class SimulationRunner:
    def __init__(self, parameters):
        self.parameters = parameters

    def runSimulation(self, runningControlFile=""):
        pass

    def hasExitCondition(self):
        return self.parameters.exitCondition is None

    def checkExitCondition(self, clustering):
        if self.parameters.exitCondition:
            return self.parameters.exitCondition.checkExitCondition(clustering)
        return False

    def makeWorkingControlFile(self, workingControlFilename, dictionary):
        inputFile = open(self.parameters.templetizedControlFile, "r")
        inputFileContent = inputFile.read()
        inputFile.close()

        inputFileTemplate = string.Template(inputFileContent)
        outputFileContent = inputFileTemplate.substitute(dictionary)

        outputFile = open(workingControlFilename, "w")
        outputFileContent = outputFileContent.replace("'", '"')
        outputFile.write(outputFileContent)
        outputFile.close()

class PeleSimulation(SimulationRunner):
    def __init__(self, parameters):
        SimulationRunner.__init__(self, parameters)
        self.type = simulationTypes.SIMULATION_TYPE.PELE

    def createSymbolicLinks(self):
        if not os.path.islink("Data"):
            os.system("ln -s " + self.parameters.dataFolder + " Data")
        if not os.path.islink("Documents"):
            os.system("ln -s " + self.parameters.documentsFolder + " Documents")

    def runSimulation(self, runningControlFile=""):
        self.createSymbolicLinks()

        toRun = ["mpirun -np " + str(self.parameters.processors), self.parameters.executable, runningControlFile]
        toRun = " ".join(toRun)
        print toRun
        startTime = time.time()
        proc = subprocess.Popen(toRun, stdout=subprocess.PIPE, shell=True)
        (out, err) = proc.communicate()
        print out
        if err:
            print err

        endTime = time.time()
        print "PELE took %.2f sec" % (endTime - startTime)

class TestSimulation(SimulationRunner):
    """
        Class used for testing
    """
    def __init__(self, parameters):
        self.type = simulationTypes.SIMULATION_TYPE.TEST
        self.copied = False
        self.parameters = parameters

    def runSimulation(self, runningControlFile=""):
        if not self.copied:
            if os.path.exists(self.parameters.destination):
                shutil.rmtree(self.parameters.destination)
            shutil.copytree(self.parameters.origin, self.parameters.destination)
            self.copied = True

    def makeWorkingControlFile(self, workingControlFilename, dictionary):
        pass

class ExitConditionBuilder:
    def build(self, exitConditionBlock):
        exitConditionType = exitConditionBlock[blockNames.ExitConditionType.type]
        exitConditionParams = exitConditionBlock[blockNames.SimulationParams.params]
        if exitConditionType == blockNames.ExitConditionType.metric:
            metricCol = exitConditionParams[blockNames.SimulationParams.metricCol]
            metricValue = exitConditionParams[blockNames.SimulationParams.exitValue]
            return MetricExitCondition(metricCol, metricValue,
                                       exitConditionType)
        else:
            sys.exit("Unknown exit condition type! Choices are: " + str(simulationTypes.EXITCONDITION_TYPE_TO_STRING_DICTIONARY.values()))

class MetricExitCondition:
    def __init__(self, metricCol, metricValue, type):
        self.metricCol = metricCol
        self.metricValue = metricValue
        self.lastCheckedCluster = 0
        self.type = simulationTypes.EXITCONDITION_TYPE.METRIC

    def checkExitCondition(self, clustering):
        """ Iterate over all unchecked cluster and check if the exit condtion
            is met
        """
        for i in range(self.lastCheckedCluster, clustering.clusters.getNumberClusters()):
            cluster = clustering.clusters.getCluster(i)
            metric = cluster.getMetricFromColumn(self.metricCol)
            if metric is not None and metric < self.metricValue:
                return True
            self.lastCheckedCluster = i
        return False

class RunnerBuilder:

    def build(self, simulationRunnerBlock):
        simulationType = simulationRunnerBlock[blockNames.SimulationType.type]
        paramsBlock = simulationRunnerBlock[blockNames.SimulationParams.params]
        params = SimulationParameters()
        if simulationType == blockNames.SimulationType.pele:
            params.processors = paramsBlock[blockNames.SimulationParams.processors]
            params.dataFolder = paramsBlock.get(blockNames.SimulationParams.dataFolder, constants.DATA_FOLDER)
            params.documentsFolder = paramsBlock.get(blockNames.SimulationParams.documentsFolder, constants.DOCUMENTS_FOLDER)
            params.executable = paramsBlock.get(blockNames.SimulationParams.executable, constants.PELE_EXECUTABLE)
            params.templetizedControlFile = paramsBlock[blockNames.SimulationParams.templetizedControlFile]
            params.iterations = paramsBlock[blockNames.SimulationParams.iterations]
            params.peleSteps = paramsBlock[blockNames.SimulationParams.peleSteps]
            params.seed = paramsBlock[blockNames.SimulationParams.seed]
            exitConditionBlock = paramsBlock.get(blockNames.SimulationParams.exitCondition, None)
            if exitConditionBlock:
                exitConditionBuilder = ExitConditionBuilder()
                params.exitCondition = exitConditionBuilder.build(exitConditionBlock)

            SimulationRunner = PeleSimulation(params)
        elif simulationType == blockNames.SimulationType.md:
            pass
        elif simulationType == blockNames.SimulationType.test:
            params.processors = paramsBlock[blockNames.SimulationParams.processors]
            params.destination = paramsBlock[blockNames.SimulationParams.destination]
            params.origin = paramsBlock[blockNames.SimulationParams.origin]
            params.iterations = paramsBlock[blockNames.SimulationParams.iterations]
            params.peleSteps = paramsBlock[blockNames.SimulationParams.peleSteps]
            params.seed = paramsBlock[blockNames.SimulationParams.seed]
            SimulationRunner = TestSimulation(params)
        else:
            sys.exit("Unknown simulation type! Choices are: " + str(simulationTypes.SIMULATION_TYPE_TO_STRING_DICTIONARY.values()))
        return SimulationRunner
<commit_msg>Fixed bug in hasExitCondition in simulationRunner<commit_after>import time
import os
from constants import constants
import subprocess
from constants import blockNames
import shutil
import simulationTypes
import string
import sys

class SimulationParameters:
    def __init__(self):
        self.processors = 0
        self.executable = ""
        self.templetizedControlFile = ""
        self.dataFolder = ""
        self.documentsFolder = ""
        self.iterations = 0
        self.peleSteps = 0
        self.seed = 0
        self.exitCondition = None

class SimulationRunner:
    def __init__(self, parameters):
        self.parameters = parameters

    def runSimulation(self, runningControlFile=""):
        pass

    def hasExitCondition(self):
        return self.parameters.exitCondition is not None

    def checkExitCondition(self, clustering):
        if self.parameters.exitCondition:
            return self.parameters.exitCondition.checkExitCondition(clustering)
        return False

    def makeWorkingControlFile(self, workingControlFilename, dictionary):
        inputFile = open(self.parameters.templetizedControlFile, "r")
        inputFileContent = inputFile.read()
        inputFile.close()

        inputFileTemplate = string.Template(inputFileContent)
        outputFileContent = inputFileTemplate.substitute(dictionary)

        outputFile = open(workingControlFilename, "w")
        outputFileContent = outputFileContent.replace("'", '"')
        outputFile.write(outputFileContent)
        outputFile.close()

class PeleSimulation(SimulationRunner):
    def __init__(self, parameters):
        SimulationRunner.__init__(self, parameters)
        self.type = simulationTypes.SIMULATION_TYPE.PELE

    def createSymbolicLinks(self):
        if not os.path.islink("Data"):
            os.system("ln -s " + self.parameters.dataFolder + " Data")
        if not os.path.islink("Documents"):
            os.system("ln -s " + self.parameters.documentsFolder + " Documents")

    def runSimulation(self, runningControlFile=""):
        self.createSymbolicLinks()

        toRun = ["mpirun -np " + str(self.parameters.processors), self.parameters.executable, runningControlFile]
        toRun = " ".join(toRun)
        print toRun
        startTime = time.time()
        proc = subprocess.Popen(toRun, stdout=subprocess.PIPE, shell=True)
        (out, err) = proc.communicate()
        print out
        if err:
            print err

        endTime = time.time()
        print "PELE took %.2f sec" % (endTime - startTime)

class TestSimulation(SimulationRunner):
    """
        Class used for testing
    """
    def __init__(self, parameters):
        self.type = simulationTypes.SIMULATION_TYPE.TEST
        self.copied = False
        self.parameters = parameters

    def runSimulation(self, runningControlFile=""):
        if not self.copied:
            if os.path.exists(self.parameters.destination):
                shutil.rmtree(self.parameters.destination)
            shutil.copytree(self.parameters.origin, self.parameters.destination)
            self.copied = True

    def makeWorkingControlFile(self, workingControlFilename, dictionary):
        pass

class ExitConditionBuilder:
    def build(self, exitConditionBlock):
        exitConditionType = exitConditionBlock[blockNames.ExitConditionType.type]
        exitConditionParams = exitConditionBlock[blockNames.SimulationParams.params]
        if exitConditionType == blockNames.ExitConditionType.metric:
            metricCol = exitConditionParams[blockNames.SimulationParams.metricCol]
            metricValue = exitConditionParams[blockNames.SimulationParams.exitValue]
            return MetricExitCondition(metricCol, metricValue,
                                       exitConditionType)
        else:
            sys.exit("Unknown exit condition type! Choices are: " + str(simulationTypes.EXITCONDITION_TYPE_TO_STRING_DICTIONARY.values()))

class MetricExitCondition:
    def __init__(self, metricCol, metricValue, type):
        self.metricCol = metricCol
        self.metricValue = metricValue
        self.lastCheckedCluster = 0
        self.type = simulationTypes.EXITCONDITION_TYPE.METRIC

    def checkExitCondition(self, clustering):
        """ Iterate over all unchecked cluster and check if the exit condtion
            is met
        """
        for i in range(self.lastCheckedCluster, clustering.clusters.getNumberClusters()):
            cluster = clustering.clusters.getCluster(i)
            metric = cluster.getMetricFromColumn(self.metricCol)
            if metric is not None and metric < self.metricValue:
                return True
            self.lastCheckedCluster = i
        return False

class RunnerBuilder:

    def build(self, simulationRunnerBlock):
        simulationType = simulationRunnerBlock[blockNames.SimulationType.type]
        paramsBlock = simulationRunnerBlock[blockNames.SimulationParams.params]
        params = SimulationParameters()
        if simulationType == blockNames.SimulationType.pele:
            params.processors = paramsBlock[blockNames.SimulationParams.processors]
            params.dataFolder = paramsBlock.get(blockNames.SimulationParams.dataFolder, constants.DATA_FOLDER)
            params.documentsFolder = paramsBlock.get(blockNames.SimulationParams.documentsFolder, constants.DOCUMENTS_FOLDER)
            params.executable = paramsBlock.get(blockNames.SimulationParams.executable, constants.PELE_EXECUTABLE)
            params.templetizedControlFile = paramsBlock[blockNames.SimulationParams.templetizedControlFile]
            params.iterations = paramsBlock[blockNames.SimulationParams.iterations]
            params.peleSteps = paramsBlock[blockNames.SimulationParams.peleSteps]
            params.seed = paramsBlock[blockNames.SimulationParams.seed]
            exitConditionBlock = paramsBlock.get(blockNames.SimulationParams.exitCondition, None)
            if exitConditionBlock:
                exitConditionBuilder = ExitConditionBuilder()
                params.exitCondition = exitConditionBuilder.build(exitConditionBlock)

            SimulationRunner = PeleSimulation(params)
        elif simulationType == blockNames.SimulationType.md:
            pass
        elif simulationType == blockNames.SimulationType.test:
            params.processors = paramsBlock[blockNames.SimulationParams.processors]
            params.destination = paramsBlock[blockNames.SimulationParams.destination]
            params.origin = paramsBlock[blockNames.SimulationParams.origin]
            params.iterations = paramsBlock[blockNames.SimulationParams.iterations]
            params.peleSteps = paramsBlock[blockNames.SimulationParams.peleSteps]
            params.seed = paramsBlock[blockNames.SimulationParams.seed]
            SimulationRunner = TestSimulation(params)
        else:
            sys.exit("Unknown simulation type! Choices are: " + str(simulationTypes.SIMULATION_TYPE_TO_STRING_DICTIONARY.values()))
        return SimulationRunner
<|endoftext|>


<commit_before># Copyright 2021 UW-IT, University of Washington
# SPDX-License-Identifier: Apache-2.0

from django.conf import settings
from rest_framework.views import APIView
from rest_framework.authentication import TokenAuthentication
from rest_framework.permissions import IsAuthenticated
from sis_provisioner.views.admin import RESTDispatch
from sis_provisioner.dao.user import (
    get_person_by_netid, get_person_by_gmail_id, can_access_canvas,
    user_sis_id, user_email)
from sis_provisioner.exceptions import UserPolicyException
from restclients_core.exceptions import DataFailureException
from logging import getLogger

logger = getLogger(__name__)

class LoginValidationView(APIView):
    authentication_classes = [TokenAuthentication]
    permission_classes = [IsAuthenticated]

    def post(self, request, *args, **kwargs):
        try:
            login_data = request.data['logins']
        except KeyError as ex:
            return RESTDispatch.error_response(400, 'Missing list of logins')

        users = []
        for login in login_data:
            login = login.strip().lower()
            if not any(u.get('login') == login for u in users):
                try:
                    user = {}
                    try:
                        person = get_person_by_gmail_id(login)
                        user['login'] = person.login_id
                        user['full_name'] = person.login_id
                    except UserPolicyException:
                        login = self.strip_domain(login)
                        person = get_person_by_netid(login)
                        user['login'] = person.uwnetid
                        try:
                            user['full_name'] = person.get_formatted_name(
                                '{first} {last}')
                        except AttributeError as ex:
                            user['full_name'] = person.display_name

                    sis_id = user_sis_id(person)
                    if not any(u.get('sis_id') == sis_id for u in users):
                        try:
                            can_access_canvas(user['login'])
                        except UserPolicyException as ex:
                            user['error'] = '{}'.format(ex)

                        user['sis_id'] = sis_id
                        user['email'] = user_email(person)
                        users.append(user)

                except DataFailureException as ex:
                    users.append({'login': login, 'error': ex.msg})

                except UserPolicyException as ex:
                    users.append({'login': login, 'error': '{}'.format(ex)})

        return RESTDispatch.json_response({'users': users})

    @staticmethod
    def strip_domain(login):
        try:
            (username, domain) = login.split('@')
            if domain in getattr(settings, 'ALLOWED_ADD_USER_DOMAINS', []):
                return username
        except ValueError:
            pass
        return login
<commit_msg>add is_person attr to login api<commit_after># Copyright 2021 UW-IT, University of Washington
# SPDX-License-Identifier: Apache-2.0

from django.conf import settings
from rest_framework.views import APIView
from rest_framework.authentication import TokenAuthentication
from rest_framework.permissions import IsAuthenticated
from sis_provisioner.views.admin import RESTDispatch
from sis_provisioner.dao.user import (
    get_person_by_netid, get_person_by_gmail_id, can_access_canvas,
    user_sis_id, user_email)
from sis_provisioner.exceptions import UserPolicyException
from restclients_core.exceptions import DataFailureException
from logging import getLogger

logger = getLogger(__name__)

class LoginValidationView(APIView):
    authentication_classes = [TokenAuthentication]
    permission_classes = [IsAuthenticated]

    def post(self, request, *args, **kwargs):
        try:
            login_data = request.data['logins']
        except KeyError as ex:
            return RESTDispatch.error_response(400, 'Missing list of logins')

        users = []
        for login in login_data:
            login = login.strip().lower()
            if not any(u.get('login') == login for u in users):
                try:
                    user = {}
                    try:
                        person = get_person_by_gmail_id(login)
                        user['login'] = person.login_id
                        user['full_name'] = person.login_id
                        user['is_person'] = True
                    except UserPolicyException:
                        login = self.strip_domain(login)
                        person = get_person_by_netid(login)
                        user['login'] = person.uwnetid
                        try:
                            user['full_name'] = person.get_formatted_name(
                                '{first} {last}')
                            user['is_person'] = True
                        except AttributeError as ex:
                            user['full_name'] = person.display_name
                            user['is_person'] = False  # UW entity

                    sis_id = user_sis_id(person)
                    if not any(u.get('sis_id') == sis_id for u in users):
                        try:
                            can_access_canvas(user['login'])
                        except UserPolicyException as ex:
                            user['error'] = '{}'.format(ex)

                        user['sis_id'] = sis_id
                        user['email'] = user_email(person)
                        users.append(user)

                except DataFailureException as ex:
                    users.append({'login': login, 'error': ex.msg})

                except UserPolicyException as ex:
                    users.append({'login': login, 'error': '{}'.format(ex)})

        return RESTDispatch.json_response({'users': users})

    @staticmethod
    def strip_domain(login):
        try:
            (username, domain) = login.split('@')
            if domain in getattr(settings, 'ALLOWED_ADD_USER_DOMAINS', []):
                return username
        except ValueError:
            pass
        return login
<|endoftext|>


<commit_before>#!/usr/bin/env python

from pathlib import Path
from multiprocessing import Pool

import nbrun

default_notebook_name = 'smFRET-Quick-Test-Server.ipynb'

def get_file_list(folder, init_filelist=None):
    folder = Path(folder)
    if init_filelist is None:
        init_filelist = []
    return [f for f in folder.glob('**/*.hdf5')
            if not f.stem.endswith('_cache')]

def run_notebook(data_filename, input_notebook=None):
    if input_notebook is None:
        input_notebook = default_notebook_name
    nbrun.run_notebook(input_notebook,
                       out_notebook_path=data_filename.with_suffix('ipynb'),
                       nb_kwargs={'fname': data_filename})

def batch_process(folder, nproc=4, notebook=None):
    assert folder.is_dir(), 'Path not found: %s' % folder

    title_msg = 'Processing files in folder: %s' % folder.name
    print('\n\n%s' % title_msg)

    filelist = get_file_list(folder)

    print('- The following files will be processed:')
    for f in filelist:
        print('  %s' % f)
    print()

    with Pool(processes=nproc) as pool:
        try:
            pool.starmap(run_notebook, [(f, notebook) for f in filelist])
        except KeyboardInterrupt:
            print('\n>>> Got keyboard interrupt.\n', flush=True)
    print('Closing subprocess pool.', flush=True)

if __name__ == '__main__':
    import argparse
    descr = """\
        This script executes an analysis notebook on all the HDF5 files
        in the passed folder (excluding files ending with '_cache').
        """
    parser = argparse.ArgumentParser(description=descr, epilog='\n')
    parser.add_argument('folder',
                        help='Source folder with files to be processed.')
    parser.add_argument('--num-processes', '-n', metavar='N', type=int, default=4,
                        help='Number of multiprocess workers to use.')
    msg = ("Filename of the analysis notebook. If not specified, the default "
           "notebook is '%s'." % default_notebook_name)
    parser.add_argument('--notebook', metavar='NB_NAME',
                        default=default_notebook_name, help=msg)
    args = parser.parse_args()

    folder = Path(args.folder)
    assert folder.is_dir(), 'Path not found: %s' % folder
    batch_process(folder, nproc=args.num_processes, notebook=args.notebook)
    print('Batch analysis completed.', flush=True)
<commit_msg>Fix typo<commit_after>#!/usr/bin/env python

from pathlib import Path
from multiprocessing import Pool

import nbrun

default_notebook_name = 'smFRET-Quick-Test-Server.ipynb'

def get_file_list(folder, init_filelist=None):
    folder = Path(folder)
    if init_filelist is None:
        init_filelist = []
    return [f for f in folder.glob('**/*.hdf5')
            if not f.stem.endswith('_cache')]

def run_notebook(data_filename, input_notebook=None):
    if input_notebook is None:
        input_notebook = default_notebook_name
    nbrun.run_notebook(input_notebook,
                       out_notebook_path=data_filename.with_suffix('.ipynb'),
                       nb_kwargs={'fname': data_filename})

def batch_process(folder, nproc=4, notebook=None):
    assert folder.is_dir(), 'Path not found: %s' % folder

    title_msg = 'Processing files in folder: %s' % folder.name
    print('\n\n%s' % title_msg)

    filelist = get_file_list(folder)

    print('- The following files will be processed:')
    for f in filelist:
        print('  %s' % f)
    print()

    with Pool(processes=nproc) as pool:
        try:
            pool.starmap(run_notebook, [(f, notebook) for f in filelist])
        except KeyboardInterrupt:
            print('\n>>> Got keyboard interrupt.\n', flush=True)
    print('Closing subprocess pool.', flush=True)

if __name__ == '__main__':
    import argparse
    descr = """\
        This script executes an analysis notebook on all the HDF5 files
        in the passed folder (excluding files ending with '_cache').
        """
    parser = argparse.ArgumentParser(description=descr, epilog='\n')
    parser.add_argument('folder',
                        help='Source folder with files to be processed.')
    parser.add_argument('--num-processes', '-n', metavar='N', type=int, default=4,
                        help='Number of multiprocess workers to use.')
    msg = ("Filename of the analysis notebook. If not specified, the default "
           "notebook is '%s'." % default_notebook_name)
    parser.add_argument('--notebook', metavar='NB_NAME',
                        default=default_notebook_name, help=msg)
    args = parser.parse_args()

    folder = Path(args.folder)
    assert folder.is_dir(), 'Path not found: %s' % folder
    batch_process(folder, nproc=args.num_processes, notebook=args.notebook)
    print('Batch analysis completed.', flush=True)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# -----------------------------------------------------------------------------
from __future__ import absolute_import, division, print_function

import numpy as np

from skbio.core.exception import RecordError
from skbio.parse.record_finder import LabeledRecordFinder

def is_fasta_label(x):
    """Checks if x looks like a FASTA label line."""
    return x.startswith('>')

def is_blank_or_comment(x):
    """Checks if x is blank or a FASTA comment line."""
    return (not x) or x.startswith('#') or x.isspace()

def is_blank(x):
    """Checks if x is blank."""
    return (not x) or x.isspace()

FastaFinder = LabeledRecordFinder(is_fasta_label, ignore=is_blank_or_comment)

def parse_fasta(infile, strict=True, label_to_name=None, finder=FastaFinder,
                label_characters='>', ignore_comment=False):
    r"""Generator of labels and sequences from a fasta file.

    Parameters
    ----------
    infile : open file object or str
        An open fasta file or a path to a fasta file.

    strict : bool
        If `True` a ``RecordError`` will be raised if no header line is found.

    label_to_name : function
        A function to apply to the sequence label (i.e., text on the header
        line) before yielding it. By default, the sequence label is returned
        with no processing. This function must take a single string as input
        and return a single string as output.

    finder : function
        The function to apply to find records in the fasta file. In general
        you should not have to change this.

    label_characters : str
        String used to indicate the beginning of a new record. In general you
        should not have to change this.

    ignore_comment : bool
        If `True`, split the sequence label on spaces, and return the label
        only as the first space separated field (i.e., the sequence
        identifier). Note: if both ``ignore_comment`` and ``label_to_name`` are
        passed, ``ignore_comment`` is ignored (both operate on the label, so
        there is potential for things to get messy otherwise).

    Returns
    -------
    two-item tuple of str
        yields the label and sequence for each entry.

    Raises
    ------
    RecordError
        If ``strict == True``, raises a ``RecordError`` if there is a fasta
        label line with no associated sequence, or a sequence with no
        associated label line (in other words, if there is a partial record).

    Examples
    --------
    Assume we have a fasta-formatted file with the following contents::

        >seq1 db-accession-149855
        CGATGTCGATCGATCGATCGATCAG
        >seq2 db-accession-34989
        CATCGATCGATCGATGCATGCATGCATG

    >>> from StringIO import StringIO
    >>> fasta_f = StringIO('>seq1 db-accession-149855\n'
    ...                    'CGATGTCGATCGATCGATCGATCAG\n'
    ...                    '>seq2 db-accession-34989\n'
    ...                    'CATCGATCGATCGATGCATGCATGCATG\n')

    We can parse this as follows:

    >>> from skbio.parse.sequences import parse_fasta
    >>> for label, seq in parse_fasta(fasta_f):
    ...     print(label, seq)
    seq1 db-accession-149855 CGATGTCGATCGATCGATCGATCAG
    seq2 db-accession-34989 CATCGATCGATCGATGCATGCATGCATG

    The sequence label or header line in a fasta file is defined as containing
    two separate pieces of information, delimited by spaces. The first space-
    separated entry is the sequence identifier, and everything following the
    first space is considered additional information (e.g., comments about the
    source of the sequence or the molecule that it encodes). Often we don't
    care about that information internally. If you want to just return the
    sequence identifier from that line, you can do this as follows:

    >>> from StringIO import StringIO
    >>> fasta_f = StringIO('>seq1 db-accession-149855\n'
    ...                    'CGATGTCGATCGATCGATCGATCAG\n'
    ...                    '>seq2 db-accession-34989\n'
    ...                    'CATCGATCGATCGATGCATGCATGCATG\n')

    >>> from skbio.parse.sequences import parse_fasta
    >>> for label, seq in parse_fasta(fasta_f, ignore_comment=True):
    ...     print(label, seq)
    seq1 CGATGTCGATCGATCGATCGATCAG
    seq2 CATCGATCGATCGATGCATGCATGCATG

    """

    for rec in finder(infile):
        # first line must be a label line
        if not rec[0][0] in label_characters:
            if strict:
                raise RecordError("Found Fasta record without label line: %s" %
                                  rec)
            else:
                continue
        # record must have at least one sequence
        if len(rec) < 2:
            if strict:
                raise RecordError("Found label line without sequences: %s" %
                                  rec)
            else:
                continue

        # remove the label character from the beginning of the label
        label = rec[0][1:].strip()
        # if the user passed a label_to_name function, apply that to the label
        if label_to_name is not None:
            label = label_to_name(label)
        # otherwise, if the user passed ignore_comment, split the label on
        # spaces, and return the first space separated field (i.e., the
        # sequence identifier)
        elif ignore_comment:
            label = label.split()[0]
        else:
            pass

        # join the sequence lines into a single string
        seq = ''.join(rec[1:])

        yield label, seq

def parse_qual(infile, full_header=False):
    r"""yields label and qual from a qual file.

    Parameters
    ----------
    infile : open file object or str
        An open fasta file or path to it.

    full_header : bool
        Return the full header or just the id

    Returns
    -------
    label : str
        The quality label
    qual : array
        The quality at each position

    Examples
    --------
    Assume we have a qual formatted file with the following contents::

        >seq1
        10 20 30 40
        >seq2
        1 2 3 4

    >>> from StringIO import StringIO
    >>> from skbio.parse.sequences import parse_qual
    >>> qual_f = StringIO('>seq1\n'
    ...                   '10 20 30 40\n'
    ...                   '>seq2\n'
    ...                   '1 2 3 4\n')
    >>> for label, qual in parse_qual(qual_f):
    ...     print(label)
    ...     print(qual)
    seq1
    [10 20 30 40]
    seq2
    [1 2 3 4]

    """
    for rec in FastaFinder(infile):
        curr_id = rec[0][1:]
        curr_qual = ' '.join(rec[1:])
        try:
            parts = np.asarray(curr_qual.split(), dtype=int)
        except ValueError:
            raise RecordError(
                "Invalid qual file. Check the format of the qual files.")
        if full_header:
            curr_pid = curr_id
        else:
            curr_pid = curr_id.split()[0]
        yield (curr_pid, parts)
<commit_msg>changed working in documentation<commit_after>#!/usr/bin/env python
# -----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# -----------------------------------------------------------------------------
from __future__ import absolute_import, division, print_function

import numpy as np

from skbio.core.exception import RecordError
from skbio.parse.record_finder import LabeledRecordFinder

def is_fasta_label(x):
    """Checks if x looks like a FASTA label line."""
    return x.startswith('>')

def is_blank_or_comment(x):
    """Checks if x is blank or a FASTA comment line."""
    return (not x) or x.startswith('#') or x.isspace()

def is_blank(x):
    """Checks if x is blank."""
    return (not x) or x.isspace()

FastaFinder = LabeledRecordFinder(is_fasta_label, ignore=is_blank_or_comment)

def parse_fasta(infile, strict=True, label_to_name=None, finder=FastaFinder,
                label_characters='>', ignore_comment=False):
    r"""Generator of labels and sequences from a fasta file.

    Parameters
    ----------
    infile : open file object or str
        An open fasta file or a path to a fasta file.

    strict : bool
        If `True` a ``RecordError`` will be raised if no header line is found.

    label_to_name : function
        A function to apply to the sequence label (i.e., text on the header
        line) before yielding it. By default, the sequence label is returned
        with no processing. This function must take a single string as input
        and return a single string as output.

    finder : function
        The function to apply to find records in the fasta file. In general
        you should not have to change this.

    label_characters : str
        String used to indicate the beginning of a new record. In general you
        should not have to change this.

    ignore_comment : bool
        If `True`, split the sequence label on spaces, and return the label
        only as the first space separated field (i.e., the sequence
        identifier). Note: if both ``ignore_comment`` and ``label_to_name`` are
        passed, ``ignore_comment`` is ignored (both operate on the label, so
        there is potential for things to get messy otherwise).

    Returns
    -------
    two-item tuple of str
        yields the label and sequence for each entry.

    Raises
    ------
    RecordError
        If ``strict == True``, raises a ``RecordError`` if there is a fasta
        label line with no associated sequence, or a sequence with no
        associated label line (in other words, if there is a partial record).

    Examples
    --------
    Assume we have a fasta-formatted file with the following contents::

        >seq1 db-accession-149855
        CGATGTCGATCGATCGATCGATCAG
        >seq2 db-accession-34989
        CATCGATCGATCGATGCATGCATGCATG

    >>> from StringIO import StringIO
    >>> fasta_f = StringIO('>seq1 db-accession-149855\n'
    ...                    'CGATGTCGATCGATCGATCGATCAG\n'
    ...                    '>seq2 db-accession-34989\n'
    ...                    'CATCGATCGATCGATGCATGCATGCATG\n')

    We can parse this as follows:

    >>> from skbio.parse.sequences import parse_fasta
    >>> for label, seq in parse_fasta(fasta_f):
    ...     print(label, seq)
    seq1 db-accession-149855 CGATGTCGATCGATCGATCGATCAG
    seq2 db-accession-34989 CATCGATCGATCGATGCATGCATGCATG

    The sequence label or header line in a fasta file is defined as containing
    two separate pieces of information, delimited by a space. The first space-
    separated entry is the sequence identifier, and everything following the
    first space is considered additional information (e.g., comments about the
    source of the sequence or the molecule that it encodes). Often we don't
    care about that information within our code. If you want to just return the
    sequence identifier from that line, you can pass ``ignore_comment=True``:

    >>> from StringIO import StringIO
    >>> fasta_f = StringIO('>seq1 db-accession-149855\n'
    ...                    'CGATGTCGATCGATCGATCGATCAG\n'
    ...                    '>seq2 db-accession-34989\n'
    ...                    'CATCGATCGATCGATGCATGCATGCATG\n')

    >>> from skbio.parse.sequences import parse_fasta
    >>> for label, seq in parse_fasta(fasta_f, ignore_comment=True):
    ...     print(label, seq)
    seq1 CGATGTCGATCGATCGATCGATCAG
    seq2 CATCGATCGATCGATGCATGCATGCATG

    """

    for rec in finder(infile):
        # first line must be a label line
        if not rec[0][0] in label_characters:
            if strict:
                raise RecordError("Found Fasta record without label line: %s" %
                                  rec)
            else:
                continue
        # record must have at least one sequence
        if len(rec) < 2:
            if strict:
                raise RecordError("Found label line without sequences: %s" %
                                  rec)
            else:
                continue

        # remove the label character from the beginning of the label
        label = rec[0][1:].strip()
        # if the user passed a label_to_name function, apply that to the label
        if label_to_name is not None:
            label = label_to_name(label)
        # otherwise, if the user passed ignore_comment, split the label on
        # spaces, and return the first space separated field (i.e., the
        # sequence identifier)
        elif ignore_comment:
            label = label.split()[0]
        else:
            pass

        # join the sequence lines into a single string
        seq = ''.join(rec[1:])

        yield label, seq

def parse_qual(infile, full_header=False):
    r"""yields label and qual from a qual file.

    Parameters
    ----------
    infile : open file object or str
        An open fasta file or path to it.

    full_header : bool
        Return the full header or just the id

    Returns
    -------
    label : str
        The quality label
    qual : array
        The quality at each position

    Examples
    --------
    Assume we have a qual formatted file with the following contents::

        >seq1
        10 20 30 40
        >seq2
        1 2 3 4

    >>> from StringIO import StringIO
    >>> from skbio.parse.sequences import parse_qual
    >>> qual_f = StringIO('>seq1\n'
    ...                   '10 20 30 40\n'
    ...                   '>seq2\n'
    ...                   '1 2 3 4\n')
    >>> for label, qual in parse_qual(qual_f):
    ...     print(label)
    ...     print(qual)
    seq1
    [10 20 30 40]
    seq2
    [1 2 3 4]

    """
    for rec in FastaFinder(infile):
        curr_id = rec[0][1:]
        curr_qual = ' '.join(rec[1:])
        try:
            parts = np.asarray(curr_qual.split(), dtype=int)
        except ValueError:
            raise RecordError(
                "Invalid qual file. Check the format of the qual files.")
        if full_header:
            curr_pid = curr_id
        else:
            curr_pid = curr_id.split()[0]
        yield (curr_pid, parts)
<|endoftext|>


<commit_before># coding: utf-8

# Author: Johannes Schönberger
#
# License: BSD 3 clause

import numpy as np

from ..base import BaseEstimator, MetaEstimatorMixin, clone
from ..utils import check_random_state, atleast2d_or_csr
from .base import LinearRegression

class RANSAC(BaseEstimator, MetaEstimatorMixin):
    """RANSAC (RANdom SAmple Consensus) algorithm.

    RANSAC is an iterative algorithm for the robust estimation of parameters
    from a subset of inliers from the complete data set. More information can
    be found in the general documentation of linear models.

    A detailed description of the algorithm can be found in the documentation
    of the ``linear_model`` sub-package.

    Parameters
    ----------
    base_estimator : object, optional
        Base estimator object which implements the following methods:

         * `fit(X, y)`: Fit model to given training data and target values.
         * `score(X)`: Returns the mean accuracy on the given test data, which
           is used for the stop criterion defined by `stop_score`.
           Additionally, the score is used to decide which of two equally
           large consensus sets is chosen as the better one.

        If base_estimator is None, then
        ``base_estimator=sklearn.linear_model.LinearRegression`` is used for
        target values of dtype float.

        Note that the current implementation only supports regression
        estimators.

    min_n_samples : int (>= 1) or float ([0, 1]), optional
        Minimum number of samples chosen randomly from original data. Treated
        as an absolute number of samples for `min_n_samples >= 1`, treated as a
        relative number `ceil(min_n_samples * X.shape[0]`) for
        `min_n_samples < 1`. By default a
        ``sklearn.linear_model.LinearRegression`` estimator is assumed and
        `min_n_samples` is chosen as ``X.shape[1] + 1``.

    residual_threshold : float, optional
        Maximum residual for a data sample to be classified as an inlier.
        By default the threshold is chosen as the MAD (median absolute
        deviation) of the target values `y`.

    is_data_valid : callable, optional
        This function is called with the randomly selected data before the
        model is fitted to it: `is_data_valid(X, y)`. If its return value is
        False the current randomly chosen sub-sample is skipped.

    is_model_valid : callable, optional
        This function is called with the estimated model and the randomly
        selected data: `is_model_valid(model, X, y)`. If its return value is
        False the current randomly chosen sub-sample is skipped.

    max_trials : int, optional
        Maximum number of iterations for random sample selection.

    stop_n_inliers : int, optional
        Stop iteration if at least this number of inliers are found.

    stop_score : float, optional
        Stop iteration if score is greater equal than this threshold.

    residual_metric : callable, optional
        Metric to reduce the dimensionality of the residuals to 1 for
        multi-dimensional target values ``y.shape[1] > 1``. By default the sum
        of absolute differences is used::

            lambda dy: np.sum(np.abs(dx), axis=1)

    random_state : integer or numpy.RandomState, optional
        The generator used to initialize the centers. If an integer is
        given, it fixes the seed. Defaults to the global numpy random
        number generator.

    Attributes
    ----------
    estimator_ : object
        Best fitted model (copy of the `base_estimator` object).

    n_trials_ : int
        Number of random selection trials until one of the stop criteria is
        met. It is always ``<= max_trials``.

    inlier_mask_ : bool array of shape [n_samples]
        Boolean mask of inliers classified as ``True``.

    References
    ----------
    .. [1] http://en.wikipedia.org/wiki/RANSAC
    .. [2] http://www.cs.columbia.edu/~belhumeur/courses/compPhoto/ransac.pdf
    .. [3] http://www.bmva.org/bmvc/2009/Papers/Paper355/Paper355.pdf
    """

    def __init__(self, base_estimator=None, min_n_samples=None,
                 residual_threshold=None, is_data_valid=None,
                 is_model_valid=None, max_trials=100,
                 stop_n_inliers=np.inf, stop_score=np.inf,
                 residual_metric=None, random_state=None):

        self.base_estimator = base_estimator
        self.min_n_samples = min_n_samples
        self.residual_threshold = residual_threshold
        self.is_data_valid = is_data_valid
        self.is_model_valid = is_model_valid
        self.max_trials = max_trials
        self.stop_n_inliers = stop_n_inliers
        self.stop_score = stop_score
        self.residual_metric = residual_metric
        self.random_state = random_state

    def fit(self, X, y):
        """Fit estimator using RANSAC algorithm.

        Parameters
        ----------
        X : array-like or sparse matrix, shape [n_samples, n_features]
            Training data.

        y : array-like, shape = [n_samples] or [n_samples, n_targets]
            Target values.

        Raises
        ------
        ValueError
            If no valid consensus set could be found. This occurs if
            `is_data_valid` and `is_model_valid` return False for all
            `max_trials` randomly chosen sub-samples.

        """
        if self.base_estimator is not None:
            base_estimator = clone(self.base_estimator)
        elif y.dtype.kind == 'f':
            base_estimator = LinearRegression()
        else:
            raise ValueError("`base_estimator` not specified.")

        if self.min_n_samples is None:
            # assume linear model by default
            min_n_samples = X.shape[1] + 1
        elif 0 < self.min_n_samples < 1:
            min_n_samples = np.ceil(self.min_n_samples * X.shape[0])
        elif self.min_n_samples >= 1:
            min_n_samples = self.min_n_samples
        else:
            raise ValueError("Value for `min_n_samples` must be scalar and "
                             "positive.")

        if self.residual_threshold is None:
            # MAD (median absolute deviation)
            residual_threshold = np.median(np.abs(y - np.median(y)))
        else:
            residual_threshold = self.residual_threshold

        if self.residual_metric is None:
            residual_metric = lambda dy: np.sum(np.abs(dy), axis=1)
        else:
            residual_metric = self.residual_metric

        random_state = check_random_state(self.random_state)

        try:  # Not all estimator accept a random_state
            base_estimator.set_params(random_state=random_state)
        except ValueError:
            pass

        n_inliers_best = 0
        score_best = np.inf
        inlier_mask_best = None
        X_inlier_best = None
        y_inlier_best = None

        # number of data samples
        n_samples = X.shape[0]
        sample_idxs = np.arange(n_samples)

        X = atleast2d_or_csr(X)
        y = np.asarray(y)

        if y.ndim == 1:
            y = y[:, None]

        for self.n_trials_ in range(1, self.max_trials + 1):

            # choose random sample set
            random_idxs = random_state.randint(0, n_samples, min_n_samples)
            X_subset = X[random_idxs]
            y_subset = y[random_idxs]

            # check if random sample set is valid
            if (self.is_data_valid is not None
                    and not self.is_data_valid(X_subset, y_subset)):
                continue

            # fit model for current random sample set
            base_estimator.fit(X_subset, y_subset)

            # check if estimated model is valid
            if (self.is_model_valid is not None and not
                    self.is_model_valid(base_estimator, X_subset,
                                        y_subset)):
                continue

            # residuals of all data for current random sample model
            residuals_subset = residual_metric(base_estimator.predict(X) - y)

            # classify data into inliers and outliers
            inlier_mask_subset = residuals_subset < residual_threshold
            n_inliers_subset = np.sum(inlier_mask_subset)

            # less inliers -> skip current random sample
            if n_inliers_subset < n_inliers_best:
                continue

            # extract inlier data set
            inlier_idxs_subset = sample_idxs[inlier_mask_subset]
            X_inlier_subset = X[inlier_idxs_subset]
            y_inlier_subset = y[inlier_idxs_subset]

            # score of inlier data set
            score_subset = base_estimator.score(X_inlier_subset,
                                                y_inlier_subset)

            # same number of inliers but worse score -> skip current random
            # sample
            if (n_inliers_subset == n_inliers_best
                    and score_subset < score_best):
                continue

            # save current random sample as best sample
            n_inliers_best = n_inliers_subset
            score_best = score_subset
            inlier_mask_best = inlier_mask_subset
            X_inlier_best = X_inlier_subset
            y_inlier_best = y_inlier_subset

            # break if sufficient number of inliers or score is reached
            if (n_inliers_best >= self.stop_n_inliers
                    or score_best >= self.stop_score):
                break

        # if none of the iterations met the required criteria
        if inlier_mask_best is None:
            raise ValueError("RANSAC could not find valid consensus set, "
                             "because `is_data_valid` and `is_model_valid` "
                             "returned False for all `max_trials` randomly "
                             "chosen sub-samples. Consider relaxing the "
                             "constraints.")

        # estimate final model using all inliers
        base_estimator.fit(X_inlier_best, y_inlier_best)

        self.estimator_ = base_estimator
        self.inlier_mask_ = inlier_mask_best

    def predict(self, X):
        """Predict using the estimated model.

        This is a wrapper for `estimator_.predict(X)`.

        Parameters
        ----------
        X : numpy array of shape [n_samples, n_features]

        Returns
        -------
        C : array, shape = [n_samples]
            Returns predicted values.
        """
        return self.estimator_.predict(X)

    def score(self, X, y):
        """Returns the score of the prediction.

        This is a wrapper for `estimator_.score(X, y)`.

        Parameters
        ----------
        X : numpy array or sparse matrix of shape [n_samples, n_features]
            Training data.

        y : numpy array of shape [n_samples, n_targets]
            Target values.

        Returns
        -------
        z : float
            Score of the prediction.
        """
        return self.estimator_.score(X, y)
<commit_msg>Fix return value description of predict<commit_after># coding: utf-8

# Author: Johannes Schönberger
#
# License: BSD 3 clause

import numpy as np

from ..base import BaseEstimator, MetaEstimatorMixin, clone
from ..utils import check_random_state, atleast2d_or_csr
from .base import LinearRegression

class RANSAC(BaseEstimator, MetaEstimatorMixin):
    """RANSAC (RANdom SAmple Consensus) algorithm.

    RANSAC is an iterative algorithm for the robust estimation of parameters
    from a subset of inliers from the complete data set. More information can
    be found in the general documentation of linear models.

    A detailed description of the algorithm can be found in the documentation
    of the ``linear_model`` sub-package.

    Parameters
    ----------
    base_estimator : object, optional
        Base estimator object which implements the following methods:

         * `fit(X, y)`: Fit model to given training data and target values.
         * `score(X)`: Returns the mean accuracy on the given test data, which
           is used for the stop criterion defined by `stop_score`.
           Additionally, the score is used to decide which of two equally
           large consensus sets is chosen as the better one.

        If base_estimator is None, then
        ``base_estimator=sklearn.linear_model.LinearRegression`` is used for
        target values of dtype float.

        Note that the current implementation only supports regression
        estimators.

    min_n_samples : int (>= 1) or float ([0, 1]), optional
        Minimum number of samples chosen randomly from original data. Treated
        as an absolute number of samples for `min_n_samples >= 1`, treated as a
        relative number `ceil(min_n_samples * X.shape[0]`) for
        `min_n_samples < 1`. By default a
        ``sklearn.linear_model.LinearRegression`` estimator is assumed and
        `min_n_samples` is chosen as ``X.shape[1] + 1``.

    residual_threshold : float, optional
        Maximum residual for a data sample to be classified as an inlier.
        By default the threshold is chosen as the MAD (median absolute
        deviation) of the target values `y`.

    is_data_valid : callable, optional
        This function is called with the randomly selected data before the
        model is fitted to it: `is_data_valid(X, y)`. If its return value is
        False the current randomly chosen sub-sample is skipped.

    is_model_valid : callable, optional
        This function is called with the estimated model and the randomly
        selected data: `is_model_valid(model, X, y)`. If its return value is
        False the current randomly chosen sub-sample is skipped.

    max_trials : int, optional
        Maximum number of iterations for random sample selection.

    stop_n_inliers : int, optional
        Stop iteration if at least this number of inliers are found.

    stop_score : float, optional
        Stop iteration if score is greater equal than this threshold.

    residual_metric : callable, optional
        Metric to reduce the dimensionality of the residuals to 1 for
        multi-dimensional target values ``y.shape[1] > 1``. By default the sum
        of absolute differences is used::

            lambda dy: np.sum(np.abs(dx), axis=1)

    random_state : integer or numpy.RandomState, optional
        The generator used to initialize the centers. If an integer is
        given, it fixes the seed. Defaults to the global numpy random
        number generator.

    Attributes
    ----------
    estimator_ : object
        Best fitted model (copy of the `base_estimator` object).

    n_trials_ : int
        Number of random selection trials until one of the stop criteria is
        met. It is always ``<= max_trials``.

    inlier_mask_ : bool array of shape [n_samples]
        Boolean mask of inliers classified as ``True``.

    References
    ----------
    .. [1] http://en.wikipedia.org/wiki/RANSAC
    .. [2] http://www.cs.columbia.edu/~belhumeur/courses/compPhoto/ransac.pdf
    .. [3] http://www.bmva.org/bmvc/2009/Papers/Paper355/Paper355.pdf
    """

    def __init__(self, base_estimator=None, min_n_samples=None,
                 residual_threshold=None, is_data_valid=None,
                 is_model_valid=None, max_trials=100,
                 stop_n_inliers=np.inf, stop_score=np.inf,
                 residual_metric=None, random_state=None):

        self.base_estimator = base_estimator
        self.min_n_samples = min_n_samples
        self.residual_threshold = residual_threshold
        self.is_data_valid = is_data_valid
        self.is_model_valid = is_model_valid
        self.max_trials = max_trials
        self.stop_n_inliers = stop_n_inliers
        self.stop_score = stop_score
        self.residual_metric = residual_metric
        self.random_state = random_state

    def fit(self, X, y):
        """Fit estimator using RANSAC algorithm.

        Parameters
        ----------
        X : array-like or sparse matrix, shape [n_samples, n_features]
            Training data.

        y : array-like, shape = [n_samples] or [n_samples, n_targets]
            Target values.

        Raises
        ------
        ValueError
            If no valid consensus set could be found. This occurs if
            `is_data_valid` and `is_model_valid` return False for all
            `max_trials` randomly chosen sub-samples.

        """
        if self.base_estimator is not None:
            base_estimator = clone(self.base_estimator)
        elif y.dtype.kind == 'f':
            base_estimator = LinearRegression()
        else:
            raise ValueError("`base_estimator` not specified.")

        if self.min_n_samples is None:
            # assume linear model by default
            min_n_samples = X.shape[1] + 1
        elif 0 < self.min_n_samples < 1:
            min_n_samples = np.ceil(self.min_n_samples * X.shape[0])
        elif self.min_n_samples >= 1:
            min_n_samples = self.min_n_samples
        else:
            raise ValueError("Value for `min_n_samples` must be scalar and "
                             "positive.")

        if self.residual_threshold is None:
            # MAD (median absolute deviation)
            residual_threshold = np.median(np.abs(y - np.median(y)))
        else:
            residual_threshold = self.residual_threshold

        if self.residual_metric is None:
            residual_metric = lambda dy: np.sum(np.abs(dy), axis=1)
        else:
            residual_metric = self.residual_metric

        random_state = check_random_state(self.random_state)

        try:  # Not all estimator accept a random_state
            base_estimator.set_params(random_state=random_state)
        except ValueError:
            pass

        n_inliers_best = 0
        score_best = np.inf
        inlier_mask_best = None
        X_inlier_best = None
        y_inlier_best = None

        # number of data samples
        n_samples = X.shape[0]
        sample_idxs = np.arange(n_samples)

        X = atleast2d_or_csr(X)
        y = np.asarray(y)

        if y.ndim == 1:
            y = y[:, None]

        for self.n_trials_ in range(1, self.max_trials + 1):

            # choose random sample set
            random_idxs = random_state.randint(0, n_samples, min_n_samples)
            X_subset = X[random_idxs]
            y_subset = y[random_idxs]

            # check if random sample set is valid
            if (self.is_data_valid is not None
                    and not self.is_data_valid(X_subset, y_subset)):
                continue

            # fit model for current random sample set
            base_estimator.fit(X_subset, y_subset)

            # check if estimated model is valid
            if (self.is_model_valid is not None and not
                    self.is_model_valid(base_estimator, X_subset,
                                        y_subset)):
                continue

            # residuals of all data for current random sample model
            residuals_subset = residual_metric(base_estimator.predict(X) - y)

            # classify data into inliers and outliers
            inlier_mask_subset = residuals_subset < residual_threshold
            n_inliers_subset = np.sum(inlier_mask_subset)

            # less inliers -> skip current random sample
            if n_inliers_subset < n_inliers_best:
                continue

            # extract inlier data set
            inlier_idxs_subset = sample_idxs[inlier_mask_subset]
            X_inlier_subset = X[inlier_idxs_subset]
            y_inlier_subset = y[inlier_idxs_subset]

            # score of inlier data set
            score_subset = base_estimator.score(X_inlier_subset,
                                                y_inlier_subset)

            # same number of inliers but worse score -> skip current random
            # sample
            if (n_inliers_subset == n_inliers_best
                    and score_subset < score_best):
                continue

            # save current random sample as best sample
            n_inliers_best = n_inliers_subset
            score_best = score_subset
            inlier_mask_best = inlier_mask_subset
            X_inlier_best = X_inlier_subset
            y_inlier_best = y_inlier_subset

            # break if sufficient number of inliers or score is reached
            if (n_inliers_best >= self.stop_n_inliers
                    or score_best >= self.stop_score):
                break

        # if none of the iterations met the required criteria
        if inlier_mask_best is None:
            raise ValueError("RANSAC could not find valid consensus set, "
                             "because `is_data_valid` and `is_model_valid` "
                             "returned False for all `max_trials` randomly "
                             "chosen sub-samples. Consider relaxing the "
                             "constraints.")

        # estimate final model using all inliers
        base_estimator.fit(X_inlier_best, y_inlier_best)

        self.estimator_ = base_estimator
        self.inlier_mask_ = inlier_mask_best

    def predict(self, X):
        """Predict using the estimated model.

        This is a wrapper for `estimator_.predict(X)`.

        Parameters
        ----------
        X : numpy array of shape [n_samples, n_features]

        Returns
        -------
        y : array, shape = [n_samples] or [n_samples, n_targets]
            Returns predicted values.
        """
        return self.estimator_.predict(X)

    def score(self, X, y):
        """Returns the score of the prediction.

        This is a wrapper for `estimator_.score(X, y)`.

        Parameters
        ----------
        X : numpy array or sparse matrix of shape [n_samples, n_features]
            Training data.

        y : numpy array of shape [n_samples, n_targets]
            Target values.

        Returns
        -------
        z : float
            Score of the prediction.
        """
        return self.estimator_.score(X, y)
<|endoftext|>