<commit_before># Copyright 2017 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import unittest

from tracing.value import heap_profiler
from tracing.value import histogram
from tracing.value import histogram_set

class HeapProfilerUnitTest(unittest.TestCase):

  def testHeapProfiler(self):
    test_data = histogram_set.HistogramSet()
    for i in xrange(10):
      test_hist = histogram.Histogram('test', 'n%')
      test_hist.AddSample(i / 10.0)
      test_data.AddHistogram(test_hist)

    histograms = heap_profiler.Profile(test_data)

    set_size_hist = histograms.GetHistogramNamed('heap:HistogramSet')
    self.assertEquals(set_size_hist.num_values, 1)
    # The exact sizes of python objects can vary between platforms and versions.
    self.assertGreater(set_size_hist.sum, 10000)

    hist_size_hist = histograms.GetHistogramNamed('heap:Histogram')
    self.assertEquals(hist_size_hist.num_values, 10)
    self.assertGreater(hist_size_hist.sum, 10000)

    related_names = hist_size_hist.diagnostics['types']
    self.assertEquals(related_names.Get('HistogramBin'), 'heap:HistogramBin')
    self.assertEquals(related_names.Get('DiagnosticMap'), 'heap:DiagnosticMap')

    properties = hist_size_hist.bins[33].diagnostic_maps[0]['properties']
    types = hist_size_hist.bins[33].diagnostic_maps[0]['types']
    self.assertEquals(len(properties), 12)
    self.assertGreater(properties.Get('_bins'), 1000)
    self.assertEquals(len(types), 4)
    self.assertGreater(types.Get('HistogramBin'), 1000)
    self.assertGreater(types.Get('(builtin types)'), 1000)

    bin_size_hist = histograms.GetHistogramNamed('heap:HistogramBin')
    self.assertEquals(bin_size_hist.num_values, 32)
    self.assertGreater(bin_size_hist.sum, 1000)

    diag_map_size_hist = histograms.GetHistogramNamed('heap:DiagnosticMap')
    self.assertEquals(diag_map_size_hist.num_values, 10)
    self.assertGreater(diag_map_size_hist.sum, 1000)

    range_size_hist = histograms.GetHistogramNamed('heap:Range')
    self.assertEquals(range_size_hist.num_values, 22)
    self.assertGreater(range_size_hist.sum, 1000)

    stats_size_hist = histograms.GetHistogramNamed('heap:RunningStatistics')
    self.assertEquals(stats_size_hist.num_values, 10)
    self.assertGreater(stats_size_hist.sum, 1000)
<commit_msg>Relax an assertion in heap_profiler_unittest.<commit_after># Copyright 2017 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import unittest

from tracing.value import heap_profiler
from tracing.value import histogram
from tracing.value import histogram_set

class HeapProfilerUnitTest(unittest.TestCase):

  def testHeapProfiler(self):
    test_data = histogram_set.HistogramSet()
    for i in xrange(10):
      test_hist = histogram.Histogram('test', 'n%')
      test_hist.AddSample(i / 10.0)
      test_data.AddHistogram(test_hist)

    histograms = heap_profiler.Profile(test_data)

    set_size_hist = histograms.GetHistogramNamed('heap:HistogramSet')
    self.assertEquals(set_size_hist.num_values, 1)
    # The exact sizes of python objects can vary between platforms and versions.
    self.assertGreater(set_size_hist.sum, 10000)

    hist_size_hist = histograms.GetHistogramNamed('heap:Histogram')
    self.assertEquals(hist_size_hist.num_values, 10)
    self.assertGreater(hist_size_hist.sum, 10000)

    related_names = hist_size_hist.diagnostics['types']
    self.assertEquals(related_names.Get('HistogramBin'), 'heap:HistogramBin')
    self.assertEquals(related_names.Get('DiagnosticMap'), 'heap:DiagnosticMap')

    properties = hist_size_hist.bins[33].diagnostic_maps[0]['properties']
    types = hist_size_hist.bins[33].diagnostic_maps[0]['types']
    self.assertGreater(len(properties), 3)
    self.assertGreater(properties.Get('_bins'), 1000)
    self.assertEquals(len(types), 4)
    self.assertGreater(types.Get('HistogramBin'), 1000)
    self.assertGreater(types.Get('(builtin types)'), 1000)

    bin_size_hist = histograms.GetHistogramNamed('heap:HistogramBin')
    self.assertEquals(bin_size_hist.num_values, 32)
    self.assertGreater(bin_size_hist.sum, 1000)

    diag_map_size_hist = histograms.GetHistogramNamed('heap:DiagnosticMap')
    self.assertEquals(diag_map_size_hist.num_values, 10)
    self.assertGreater(diag_map_size_hist.sum, 1000)

    range_size_hist = histograms.GetHistogramNamed('heap:Range')
    self.assertEquals(range_size_hist.num_values, 22)
    self.assertGreater(range_size_hist.sum, 1000)

    stats_size_hist = histograms.GetHistogramNamed('heap:RunningStatistics')
    self.assertEquals(stats_size_hist.num_values, 10)
    self.assertGreater(stats_size_hist.sum, 1000)
<|endoftext|>


<commit_before>import utilities
import subprocess
import datetime
import numpy as np
import rasterio

def forest_age_category(tile_id):

    upload_dir = 's3://gfw2-data/climate/carbon_model/forest_age_category/20180912/'

    print "Processing:", tile_id

    ymax, xmin, ymin, xmax = utilities.coords(tile_id)

    tropics = 0

    if ymax > -30 & ymax <= 30:

        tropics = 1

    # start time
    start = datetime.datetime.now()

    # Names of the loss, gain and tree cover density tiles
    loss = '{}.tif'.format(tile_id)
    gain = 'Hansen_GFC2015_gain_{}.tif'.format(tile_id)
    tcd = 'Hansen_GFC2014_treecover2000_{}.tif'.format(tile_id)
    ifl = '{}_res_ifl_2000tif'.format(tile_id)
    biomass = '{}_biomass.tif'.format(tile_id)
    cont_eco = 'fao_ecozones_{0}.tif'.format(tile_id)

    print "  Reading input files and evaluating condtions"

    # open loss and grab metadata about the tif,
    # like its location / projection / cellsize
    with rasterio.open(loss) as loss_src:
        kwargs = loss_src.meta

        # also grab the windows (stripes) so we can iterate
        # over the entire tif without running out of memory
        windows = loss_src.block_windows(1)

        # open gain
        with rasterio.open(gain) as gain_src:
            # open extent
            with rasterio.open(tcd) as extent_src:
                # update those kwargs for the dataset we'll write out
                kwargs.update(
                    driver='GTiff',
                    count=1,
                    compress='lzw',
                    nodata=0
                )

                with rasterio.open('forest_age_category_{}.tif'.format(tile_id), 'w', **kwargs) as dst:
                    for idx, window in windows:
                        loss = loss_src.read(1, window=window)
                        gain = gain_src.read(1, window=window)
                        tcd = extent_src.read(1, window=window)

                        # create an empty array
                        dst_data = np.zeros((window.height, window.width), dtype='uint8')

                        # where loss & gain, set output to 100, otherwise keep dst_data value
                        dst_data[np.where((tcd > 0) & (loss >= 0) & (gain >= 0))] = 100

                        dst_data[np.where((tcd > 0) & (gain == 0) & (loss == 0) & (tropics == 0))] = 1
                        dst_data[np.where((tcd > 0) & (gain == 0) & (loss == 0) & (tropics == 1) & (ifl == 0))] = 2
                        dst_data[np.where((tcd > 0) & (gain == 0) & (loss == 0) & (tropics == 1) & (ifl == 1))] = 3

                        dst_data[np.where((tcd > 0) & (gain_data == 0) & (loss_data > 0) & (ifl == 1))] = 6

                        # # where loss & gain, set output to 100, otherwise keep dst_data value
                        # dst_data[np.where((loss_data > 0) & (gain_data > 0))] = 100
                        #
                        # # etc
                        # dst_data[np.where((loss_data > 10) & (gain_data > 0) & (extent_data > 50))] = 99
                        # dst_data[np.where((loss_data < 10) & (gain_data == 0) & (extent_data < 50))] = 1
                        # dst_data[np.where((loss_data > 10) & (gain_data > 0) & (extent_data >= 90))] = 2

                        dst.write_band(1, dst_data, window=window)

    pattern = 'forest_age_category'

    utilities.upload_final(pattern, upload_dir, tile_id)

    end = datetime.datetime.now()
    elapsed_time = end-start

    print "Processing time for tile", tile_id, ":", elapsed_time

<commit_msg>Working on script to create forest age category raster using rasterio.<commit_after>import utilities
import subprocess
import datetime
import numpy as np
import rasterio

def forest_age_category(tile_id):

    upload_dir = 's3://gfw2-data/climate/carbon_model/forest_age_category/20180912/'

    print "Processing:", tile_id

    ymax, xmin, ymin, xmax = utilities.coords(tile_id)

    tropics = 0

    if ymax > -30 & ymax <= 30:

        tropics = 1

    # start time
    start = datetime.datetime.now()

    # Names of the loss, gain and tree cover density tiles
    loss = '{}.tif'.format(tile_id)
    gain = 'Hansen_GFC2015_gain_{}.tif'.format(tile_id)
    tcd = 'Hansen_GFC2014_treecover2000_{}.tif'.format(tile_id)
    ifl = '{}_res_ifl_2000tif'.format(tile_id)
    biomass = '{}_biomass.tif'.format(tile_id)
    cont_eco = 'fao_ecozones_{0}.tif'.format(tile_id)

    print "  Reading input files and evaluating condtions"

    # open loss and grab metadata about the tif,
    # like its location / projection / cellsize
    with rasterio.open(loss) as loss_src:
        kwargs = loss_src.meta

        # also grab the windows (stripes) so we can iterate
        # over the entire tif without running out of memory
        windows = loss_src.block_windows(1)

        # open gain
        with rasterio.open(gain) as gain_src:
            # open extent
            with rasterio.open(tcd) as extent_src:
                # update those kwargs for the dataset we'll write out
                kwargs.update(
                    driver='GTiff',
                    count=1,
                    compress='lzw',
                    nodata=0
                )

                with rasterio.open('forest_age_category_{}.tif'.format(tile_id), 'w', **kwargs) as dst:
                    for idx, window in windows:
                        loss = loss_src.read(1, window=window)
                        gain = gain_src.read(1, window=window)
                        tcd = extent_src.read(1, window=window)

                        # create an empty array
                        dst_data = np.zeros((window.height, window.width), dtype='uint8')

                        # where loss & gain, set output to 100, otherwise keep dst_data value
                        dst_data[np.where((tcd > 0) & (loss >= 0) & (gain >= 0))] = 100

                        dst_data[np.where((tcd > 0) & (gain == 0) & (loss == 0) & (tropics == 0))] = 1
                        dst_data[np.where((tcd > 0) & (gain == 0) & (loss == 0) & (tropics == 1) & (ifl == 0))] = 2
                        dst_data[np.where((tcd > 0) & (gain == 0) & (loss == 0) & (tropics == 1) & (ifl == 1))] = 3

                        dst_data[np.where((tcd > 0) & (gain == 0) & (loss > 0) & (ifl == 1))] = 6

                        # # where loss & gain, set output to 100, otherwise keep dst_data value
                        # dst_data[np.where((loss_data > 0) & (gain_data > 0))] = 100
                        #
                        # # etc
                        # dst_data[np.where((loss_data > 10) & (gain_data > 0) & (extent_data > 50))] = 99
                        # dst_data[np.where((loss_data < 10) & (gain_data == 0) & (extent_data < 50))] = 1
                        # dst_data[np.where((loss_data > 10) & (gain_data > 0) & (extent_data >= 90))] = 2

                        dst.write_band(1, dst_data, window=window)

    pattern = 'forest_age_category'

    utilities.upload_final(pattern, upload_dir, tile_id)

    end = datetime.datetime.now()
    elapsed_time = end-start

    print "Processing time for tile", tile_id, ":", elapsed_time

<|endoftext|>


<commit_before># coding=utf-8

"""
Collects Cassandra JMX metrics from the Jolokia Agent.  Extends the
JolokiaCollector to interpret Histogram beans with information about the
distribution of request latencies.

#### Example Configuration
CassandraJolokiaCollector uses a regular expression to determine which
attributes represent histograms. This regex can be overridden by providing a
`histogram_regex` in your configuration.  You can also override `percentiles` to
collect specific percentiles from the histogram statistics.  The format is shown
below with the default values.

CassandraJolokiaCollector.conf

```
    percentiles '50,95,99'
    histogram_regex '.*HistogramMicros$'
```
"""

from diamond.collector import str_to_bool
from jolokia import JolokiaCollector
import math
import string
import re

class CassandraJolokiaCollector(JolokiaCollector):
    # override to allow setting which percentiles will be collected
    def get_default_config_help(self):
        config_help = super(CassandraJolokiaCollector,
                            self).get_default_config_help()
        config_help.update({
            'percentiles':
            'Comma separated list of percentiles to be collected '
            '(e.g., "50,95,99").',
            'histogram_regex':
            'Filter to only process attributes that match this regex',
            'nested': 'Whether or not to enable nested values from JMX'
        })
        return config_help

    # override to allow setting which percentiles will be collected
    def get_default_config(self):
        config = super(CassandraJolokiaCollector, self).get_default_config()
        config.update({
            'percentiles': ['50', '95', '99'],
            'histogram_regex': '.*HistogramMicros$',
            'nested': 'False'
        })
        return config

    def __init__(self, *args, **kwargs):
        super(CassandraJolokiaCollector, self).__init__(*args, **kwargs)
        self.offsets = self.create_offsets(91)
        self.update_config(self.config)

    def update_config(self, config):
        if 'percentiles' in config:
            self.percentiles = map(int, config['percentiles'])
        if 'histogram_regex' in config:
            self.histogram_regex = re.compile(config['histogram_regex'])

    def collect_bean(self, prefix, obj):
        for k, v in obj.iteritems():
            if type(v) in [int, float, long]:
                self.parse_and_publish(prefix, k, v)
            elif type(v) in [dict] and str_to_bool(self.config['nested']):
                self.collect_bean("%s.%s" % (prefix, k), v)
            elif type(v) in [list] and str_to_bool(self.config['nested']):
                self.interpret_bean_with_list("%s.%s" % (prefix, k), v)

    def parse_and_publish(self, prefix, key, value):
        metric_prefix, meta = prefix.split(':', 2)
        name, metric_type, self.dimensions = self.parse_meta(meta)

        metric_name_list = [metric_prefix]
        if self.config.get('prefix', None):
            metric_name_list = [self.config['prefix'], metric_prefix]
        if metric_type:
            metric_name_list.append(metric_type)
        if name:
            metric_name_list.append(name)
        if key.lower() != 'value':
            metric_name_list.append(key.lower())

        metric_name = '.'.join(metric_name_list)
        metric_name = self.clean_up(metric_name)
        if metric_name == "":
            self.dimensions = {}
            return

        if key.lower() == 'count':
            self.publish_cumulative_counter(metric_name, value)
        else:
            self.publish(metric_name, value)

    def parse_meta(self, meta):
        dimensions = {}
        for k, v in [kv.split('=') for kv in meta.split(',')]:
            dimensions[str(k)] = v

        # We can have no name defined for metrics like:
        # org.apache.cassandra.auth:type=PermissionsCache
        metric_name = dimensions.pop("name", None)
        metric_type = dimensions.pop("type", None)
        scope_type = dimensions.pop("scope", None)
        if scope_type:
            dimensions["type"] = scope_type

        return metric_name, metric_type, dimensions

    # override: Interpret beans that match the `histogram_regex` as histograms,
    # and collect percentiles from them.
    def interpret_bean_with_list(self, prefix, values):
        if not self.histogram_regex.match(prefix):
            return

        buckets = values
        offsets = self.offsets
        for percentile in self.percentiles:
            value = self.compute_percentile(offsets, buckets, percentile)
            cleaned_key = self.clean_up("%s.p%s" % (prefix, percentile))
            self.publish(cleaned_key, value)

    # Adapted from Cassandra docs:
    # https://bit.ly/13M5JPE
    # The index corresponds to the x-axis in a histogram.  It represents buckets
    # of values, which are a series of ranges. Each offset includes the range of
    # values greater than the previous offset and less than or equal to the
    # current offset. The offsets start at 1 and each subsequent offset is
    # calculated by multiplying the previous offset by 1.2, rounding up, and
    # removing duplicates. The offsets can range from 1 to approximately 25
    # million, with less precision as the offsets get larger.
    def compute_percentile(self, offsets, buckets, percentile_int):
        non_zero_points_sum = sum(buckets)
        if non_zero_points_sum is 0:
            return 0
        middle_point_index = math.floor(
            non_zero_points_sum * (percentile_int / float(100)))

        points_seen = 0
        for index, bucket in enumerate(buckets):
            points_seen += bucket
            if points_seen >= middle_point_index:
                return round((offsets[index] - offsets[index - 1]) / 2)

    # Returns a list of offsets for `n` buckets.
    def create_offsets(self, bucket_count):
        last_num = 1
        offsets = [last_num]

        for index in range(bucket_count):
            next_num = round(last_num * 1.2)
            if next_num == last_num:
                next_num += 1
            offsets.append(next_num)
            last_num = next_num

        return offsets
<commit_msg>Using isinstance instead of type<commit_after># coding=utf-8

"""
Collects Cassandra JMX metrics from the Jolokia Agent.  Extends the
JolokiaCollector to interpret Histogram beans with information about the
distribution of request latencies.

#### Example Configuration
CassandraJolokiaCollector uses a regular expression to determine which
attributes represent histograms. This regex can be overridden by providing a
`histogram_regex` in your configuration.  You can also override `percentiles` to
collect specific percentiles from the histogram statistics.  The format is shown
below with the default values.

CassandraJolokiaCollector.conf

```
    percentiles '50,95,99'
    histogram_regex '.*HistogramMicros$'
```
"""

from diamond.collector import str_to_bool
from jolokia import JolokiaCollector
import math
import string
import re

class CassandraJolokiaCollector(JolokiaCollector):
    # override to allow setting which percentiles will be collected
    def get_default_config_help(self):
        config_help = super(CassandraJolokiaCollector,
                            self).get_default_config_help()
        config_help.update({
            'percentiles':
            'Comma separated list of percentiles to be collected '
            '(e.g., "50,95,99").',
            'histogram_regex':
            'Filter to only process attributes that match this regex',
            'nested': 'Whether or not to enable nested values from JMX'
        })
        return config_help

    # override to allow setting which percentiles will be collected
    def get_default_config(self):
        config = super(CassandraJolokiaCollector, self).get_default_config()
        config.update({
            'percentiles': ['50', '95', '99'],
            'histogram_regex': '.*HistogramMicros$',
            'nested': 'False'
        })
        return config

    def __init__(self, *args, **kwargs):
        super(CassandraJolokiaCollector, self).__init__(*args, **kwargs)
        self.offsets = self.create_offsets(91)
        self.update_config(self.config)

    def update_config(self, config):
        if 'percentiles' in config:
            self.percentiles = map(int, config['percentiles'])
        if 'histogram_regex' in config:
            self.histogram_regex = re.compile(config['histogram_regex'])

    def collect_bean(self, prefix, obj):
        for k, v in obj.iteritems():
            if isinstance(v, (int, float, long)):
                self.parse_and_publish(prefix, k, v)
            elif isinstance(v, dict) and str_to_bool(self.config['nested']):
                self.collect_bean("%s.%s" % (prefix, k), v)
            elif isinstance(v, list) and str_to_bool(self.config['nested']):
                self.interpret_bean_with_list("%s.%s" % (prefix, k), v)

    def parse_and_publish(self, prefix, key, value):
        metric_prefix, meta = prefix.split(':', 2)
        name, metric_type, self.dimensions = self.parse_meta(meta)

        metric_name_list = [metric_prefix]
        if self.config.get('prefix', None):
            metric_name_list = [self.config['prefix'], metric_prefix]
        if metric_type:
            metric_name_list.append(metric_type)
        if name:
            metric_name_list.append(name)
        if key.lower() != 'value':
            metric_name_list.append(key.lower())

        metric_name = '.'.join(metric_name_list)
        metric_name = self.clean_up(metric_name)
        if metric_name == "":
            self.dimensions = {}
            return

        if key.lower() == 'count':
            self.publish_cumulative_counter(metric_name, value)
        else:
            self.publish(metric_name, value)

    def parse_meta(self, meta):
        dimensions = {}
        for k, v in [kv.split('=') for kv in meta.split(',')]:
            dimensions[str(k)] = v

        # We can have no name defined for metrics like:
        # org.apache.cassandra.auth:type=PermissionsCache
        metric_name = dimensions.pop("name", None)
        metric_type = dimensions.pop("type", None)
        scope_type = dimensions.pop("scope", None)
        if scope_type:
            dimensions["type"] = scope_type

        return metric_name, metric_type, dimensions

    # override: Interpret beans that match the `histogram_regex` as histograms,
    # and collect percentiles from them.
    def interpret_bean_with_list(self, prefix, values):
        if not self.histogram_regex.match(prefix):
            return

        buckets = values
        offsets = self.offsets
        for percentile in self.percentiles:
            value = self.compute_percentile(offsets, buckets, percentile)
            cleaned_key = self.clean_up("%s.p%s" % (prefix, percentile))
            self.publish(cleaned_key, value)

    # Adapted from Cassandra docs:
    # https://bit.ly/13M5JPE
    # The index corresponds to the x-axis in a histogram.  It represents buckets
    # of values, which are a series of ranges. Each offset includes the range of
    # values greater than the previous offset and less than or equal to the
    # current offset. The offsets start at 1 and each subsequent offset is
    # calculated by multiplying the previous offset by 1.2, rounding up, and
    # removing duplicates. The offsets can range from 1 to approximately 25
    # million, with less precision as the offsets get larger.
    def compute_percentile(self, offsets, buckets, percentile_int):
        non_zero_points_sum = sum(buckets)
        if non_zero_points_sum is 0:
            return 0
        middle_point_index = math.floor(
            non_zero_points_sum * (percentile_int / float(100)))

        points_seen = 0
        for index, bucket in enumerate(buckets):
            points_seen += bucket
            if points_seen >= middle_point_index:
                return round((offsets[index] - offsets[index - 1]) / 2)

    # Returns a list of offsets for `n` buckets.
    def create_offsets(self, bucket_count):
        last_num = 1
        offsets = [last_num]

        for index in range(bucket_count):
            next_num = round(last_num * 1.2)
            if next_num == last_num:
                next_num += 1
            offsets.append(next_num)
            last_num = next_num

        return offsets
<|endoftext|>


<commit_before>'''
File server pluggable modules and generic backend functions
'''

# Import python libs
import os
import re
import fnmatch
import logging

# Import salt libs
import salt.loader

log = logging.getLogger(__name__)

def is_file_ignored(opts, fname):
    '''
    If file_ignore_regex or file_ignore_glob were given in config,
    compare the given file path against all of them and return True
    on the first match.
    '''
    if opts['file_ignore_regex']:
        for regex in opts['file_ignore_regex']:
            if re.search(regex, fname):
                log.debug(
                    'File matching file_ignore_regex. Skipping: {0}'.format(
                        fname
                    )
                )
                return True

    if opts['file_ignore_glob']:
        for glob in opts['file_ignore_glob']:
            if fnmatch.fnmatch(fname, glob):
                log.debug(
                    'File matching file_ignore_glob. Skipping: {0}'.format(
                        fname
                    )
                )
                return True
    return False

class Fileserver(object):
    '''
    Create a fileserver wrapper object that wraps the fileserver functions and
    iterates over them to execute the desired function within the scope of the
    desired fileserver backend.
    '''
    def __init__(self, opts):
        self.opts = opts
        self.servers = salt.loader.fileserver(opts, opts['fileserver_backend'])

    def _gen_back(self, back):
        '''
        Return the backend list
        '''
        if not back:
            back = self.opts['fileserver_backend']
        if isinstance(back, str):
            back = [back]
        return back

    def update(self, back=None):
        '''
        Update all of the fileservers that support the update function or the
        named fileserver only.
        '''
        back = self._gen_back(back)
        for fsb in back:
            fstr = '{0}.update'.format(fsb)
            if fstr in self.servers:
                self.servers[fstr]()

    def envs(self, back=None, sources=False):
        '''
        Return the environments for the named backend or all backends
        '''
        back = self._gen_back(back)
        ret = set()
        if sources:
            ret = {}
        for fsb in back:
            fstr = '{0}.envs'.format(fsb)
            if sources:
                ret[fsb] = self.servers[fstr]()
            else:
                ret.update(self.servers[fstr]())
        if sources:
            return ret
        return list(ret)

    def init(self, back=None):
        '''
        Initialize the backend, only do so if the fs supports an init function
        '''
        back = self._gen_back(back)
        for fsb in back:
            fstr = '{0}.init'.format(fsb)
            if fstr in self.servers:
                self.servers[fstr]()

    def find_file(self, path, env, back=None):
        '''
        Find the path and return the fnd structure, this structure is passed
        to other backend interfaces.
        '''
        back = self._gen_back(back)
        kwargs = {}
        fnd = {'path': '',
               'rel': ''}
        if os.path.isabs(path):
            return fnd
        if path.startswith('|'):
            # The path arguments are escaped
            path = path[1:]
        else:
            if '?' in path:
                hcomps = path.split('?')
                path = hcomps[0]
                comps = hcomps[1].split('&')
                for comp in comps:
                    if not '=' in comp:
                        # Invalid option, skip it
                        continue
                    args = comp.split('=', 1)
                    kwargs[comp[0]] = comp[1]
        if 'env' in kwargs:
            env = kwargs.pop('env')
        for fsb in back:
            fstr = '{0}.find_file'.format(fsb)
            if fstr in self.servers:
                fnd = self.servers[fstr](path, env, **kwargs)
                if fnd.get('path'):
                    fnd['back'] = fsb
                    return fnd
        return fnd

    def serve_file(self, load):
        '''
        Serve up a chunk of a file
        '''
        ret = {'data': '',
               'dest': ''}
        if 'path' not in load or 'loc' not in load or 'env' not in load:
            return ret
        fnd = self.find_file(load['path'], load['env'])
        if not fnd.get('back'):
            return ret
        fstr = '{0}.serve_file'.format(fnd['back'])
        if fstr in self.servers:
            return self.servers[fstr](load, fnd)
        return ret

    def file_hash(self, load):
        '''
        Return the hash of a given file
        '''
        if 'path' not in load or 'env' not in load:
            return ''
        fnd = self.find_file(load['path'], load['env'])
        if not fnd.get('back'):
            return ''
        fstr = '{0}.file_hash'.format(fnd['back'])
        if fstr in self.servers:
            return self.servers[fstr](load, fnd)
        return ''

    def file_list(self, load):
        '''
        Return a list of files from the dominant environment
        '''
        ret = set()
        if 'env' not in load:
            return []
        for fsb in self._gen_back(None):
            fstr = '{0}.file_list'.format(fsb)
            if fstr in self.servers:
                ret.update(self.servers[fstr](load))
        return sorted(ret)

    def file_list_emptydirs(self, load):
        '''
        List all emptydirs in the given environment
        '''
        ret = set()
        if 'env' not in load:
            return []
        for fsb in self._gen_back(None):
            fstr = '{0}.file_list_emptydirs'.format(fsb)
            if fstr in self.servers:
                ret.update(self.servers[fstr](load))
        return sorted(ret)

    def dir_list(self, load):
        '''
        List all directories in the given environment
        '''
        ret = set()
        if 'env' not in load:
            return []
        for fsb in self._gen_back(None):
            fstr = '{0}.dir_list'.format(fsb)
            if fstr in self.servers:
                ret.update(self.servers[fstr](load))
        return sorted(ret)
<commit_msg>Fix final issue in gitfs index support<commit_after>'''
File server pluggable modules and generic backend functions
'''

# Import python libs
import os
import re
import fnmatch
import logging

# Import salt libs
import salt.loader

log = logging.getLogger(__name__)

def is_file_ignored(opts, fname):
    '''
    If file_ignore_regex or file_ignore_glob were given in config,
    compare the given file path against all of them and return True
    on the first match.
    '''
    if opts['file_ignore_regex']:
        for regex in opts['file_ignore_regex']:
            if re.search(regex, fname):
                log.debug(
                    'File matching file_ignore_regex. Skipping: {0}'.format(
                        fname
                    )
                )
                return True

    if opts['file_ignore_glob']:
        for glob in opts['file_ignore_glob']:
            if fnmatch.fnmatch(fname, glob):
                log.debug(
                    'File matching file_ignore_glob. Skipping: {0}'.format(
                        fname
                    )
                )
                return True
    return False

class Fileserver(object):
    '''
    Create a fileserver wrapper object that wraps the fileserver functions and
    iterates over them to execute the desired function within the scope of the
    desired fileserver backend.
    '''
    def __init__(self, opts):
        self.opts = opts
        self.servers = salt.loader.fileserver(opts, opts['fileserver_backend'])

    def _gen_back(self, back):
        '''
        Return the backend list
        '''
        if not back:
            back = self.opts['fileserver_backend']
        if isinstance(back, str):
            back = [back]
        return back

    def update(self, back=None):
        '''
        Update all of the fileservers that support the update function or the
        named fileserver only.
        '''
        back = self._gen_back(back)
        for fsb in back:
            fstr = '{0}.update'.format(fsb)
            if fstr in self.servers:
                self.servers[fstr]()

    def envs(self, back=None, sources=False):
        '''
        Return the environments for the named backend or all backends
        '''
        back = self._gen_back(back)
        ret = set()
        if sources:
            ret = {}
        for fsb in back:
            fstr = '{0}.envs'.format(fsb)
            if sources:
                ret[fsb] = self.servers[fstr]()
            else:
                ret.update(self.servers[fstr]())
        if sources:
            return ret
        return list(ret)

    def init(self, back=None):
        '''
        Initialize the backend, only do so if the fs supports an init function
        '''
        back = self._gen_back(back)
        for fsb in back:
            fstr = '{0}.init'.format(fsb)
            if fstr in self.servers:
                self.servers[fstr]()

    def find_file(self, path, env, back=None):
        '''
        Find the path and return the fnd structure, this structure is passed
        to other backend interfaces.
        '''
        back = self._gen_back(back)
        kwargs = {}
        fnd = {'path': '',
               'rel': ''}
        if os.path.isabs(path):
            return fnd
        if path.startswith('|'):
            # The path arguments are escaped
            path = path[1:]
        else:
            if '?' in path:
                hcomps = path.split('?')
                path = hcomps[0]
                comps = hcomps[1].split('&')
                for comp in comps:
                    if not '=' in comp:
                        # Invalid option, skip it
                        continue
                    args = comp.split('=', 1)
                    kwargs[args[0]] = args[1]
        if 'env' in kwargs:
            env = kwargs.pop('env')
        for fsb in back:
            fstr = '{0}.find_file'.format(fsb)
            if fstr in self.servers:
                fnd = self.servers[fstr](path, env, **kwargs)
                if fnd.get('path'):
                    fnd['back'] = fsb
                    return fnd
        return fnd

    def serve_file(self, load):
        '''
        Serve up a chunk of a file
        '''
        ret = {'data': '',
               'dest': ''}
        if 'path' not in load or 'loc' not in load or 'env' not in load:
            return ret
        fnd = self.find_file(load['path'], load['env'])
        if not fnd.get('back'):
            return ret
        fstr = '{0}.serve_file'.format(fnd['back'])
        if fstr in self.servers:
            return self.servers[fstr](load, fnd)
        return ret

    def file_hash(self, load):
        '''
        Return the hash of a given file
        '''
        if 'path' not in load or 'env' not in load:
            return ''
        fnd = self.find_file(load['path'], load['env'])
        if not fnd.get('back'):
            return ''
        fstr = '{0}.file_hash'.format(fnd['back'])
        if fstr in self.servers:
            return self.servers[fstr](load, fnd)
        return ''

    def file_list(self, load):
        '''
        Return a list of files from the dominant environment
        '''
        ret = set()
        if 'env' not in load:
            return []
        for fsb in self._gen_back(None):
            fstr = '{0}.file_list'.format(fsb)
            if fstr in self.servers:
                ret.update(self.servers[fstr](load))
        return sorted(ret)

    def file_list_emptydirs(self, load):
        '''
        List all emptydirs in the given environment
        '''
        ret = set()
        if 'env' not in load:
            return []
        for fsb in self._gen_back(None):
            fstr = '{0}.file_list_emptydirs'.format(fsb)
            if fstr in self.servers:
                ret.update(self.servers[fstr](load))
        return sorted(ret)

    def dir_list(self, load):
        '''
        List all directories in the given environment
        '''
        ret = set()
        if 'env' not in load:
            return []
        for fsb in self._gen_back(None):
            fstr = '{0}.dir_list'.format(fsb)
            if fstr in self.servers:
                ret.update(self.servers[fstr](load))
        return sorted(ret)
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""Script to load project info from Lims into the project database in statusdb.

Maya Brandi, Science for Life Laboratory, Stockholm, Sweden.
"""
from __future__ import print_function
from genologics.config import BASEURI, USERNAME, PASSWORD
from genologics.lims import *
from LIMS2DB.objectsDB.functions import *
from optparse import OptionParser
from statusdb.db.utils import *
from LIMS2DB.utils import formatStack

from pprint import pprint

import codecs
import datetime
import LIMS2DB.objectsDB.objectsDB as DB
import logging
import logging.handlers
import multiprocessing as mp
import os
import Queue
import sys
import time
import traceback

class PSUL():
    def __init__(self, proj, samp_db, proj_db, upload_data, man_name, output_f, log):
        self.proj = proj
        self.id = proj.id
        self.udfs = proj.udf
        self.name = proj.name
        self.open_date = proj.open_date
        self.close_date = proj.close_date
        self.samp_db = samp_db
        self.proj_db = proj_db
        self.upload_data = upload_data
        self.man_name = man_name
        self.output_f = output_f
        self.ordered_opened = self.get_ordered_opened()
        self.lims = Lims(BASEURI, USERNAME, PASSWORD)
        self.log=log

    def print_couchdb_obj_to_file(self, obj):
        if self.output_f is not None:
            with open(self.output_f, 'w') as f:
                print(obj, file = f)
        else:
            print(obj, file = sys.stdout)

    def get_ordered_opened(self):
        """Is project registered as opened or ordered?"""

        if self.open_date:
            return self.open_date
        elif 'Order received' in self.proj.udfs:
            return self.udfs['Order received'].isoformat()
        else:
            return datetime.date.today().strftime("%Y-%m-%d")

    def determine_update(self):
        """Determine wether to and how to update project"""
        opened_after_130630 = comp_dates('2013-06-30', self.ordered_opened)
        log_info = ''
        if not opened_after_130630 :
                start_update = False
                self.log.info('Project is not updated because: '
                    'It was opened or ordered before 2013-06-30 : '
                    '({ord_op})'.format(ord_op = self.ordered_opened))
        else:
            start_update = True

        return start_update

    def update_project(self):
        """Fetch project info and update project in the database."""
        self.log.info('Handling {proj}'.format(proj = self.name))
        project = DB.ProjectDB(self.lims, self.id, self.samp_db, self.log)

        key = find_proj_from_view(self.proj_db, self.name)
        project.obj['_id'] = find_or_make_key(key)
        if self.upload_data:
            info = save_couchdb_obj(self.proj_db, project.obj)
        else:
            info = self.print_couchdb_obj_to_file(project.obj)

        self.log.info("project {name} is handled and {info}: _id = {id}".format(
                           name=self.name, info=info, id=project.obj['_id']))

    def handle_project(self):
        if self.determine_update():
            self.update_project()

def main(options):
    conf = options.conf
    upload_data = options.upload
    output_f = options.output_f
    couch = load_couch_server(conf)
    proj_db = couch['projects']
    samp_db = couch['samples']
    mainlims = Lims(BASEURI, USERNAME, PASSWORD)
    mainlog = logging.getLogger('psullogger')
    mainlog.setLevel(level=logging.INFO)
    mfh = logging.handlers.RotatingFileHandler(options.logfile, maxBytes=209715200, backupCount=5)
    mft = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    mfh.setFormatter(mft)
    mainlog.addHandler(mfh)

    if options.all_projects:
        if options.hours:
            delta=datetime.timedelta(hours=options.hours)
            time_string=(datetime.datetime.now()-delta).strftime('%Y-%m-%dT%H:%M:%SCET')
            projects=mainlims.get_projects(last_modified=time_string)
            mainlog.info("project list : {0}".format(" ".join([p.id for p in projects])))
        else:
            projects = mainlims.get_projects()
        masterProcess(options,projects, mainlims, mainlog)
    elif options.project_name:
        proj = mainlims.get_projects(name = options.project_name)
        if not proj:
            mainlog.warn('No project named {man_name} in Lims'.format(
                        man_name = options.project_name))
        else:
            P = PSUL(proj[0], samp_db, proj_db, options.upload, options.project_name, output_f, mainlog)
            P.handle_project()

def processPSUL(options, queue, logqueue):
    couch = load_couch_server(options.conf)
    proj_db = couch['projects']
    samp_db = couch['samples']
    mylims = Lims(BASEURI, USERNAME, PASSWORD)
    work=True
    procName=mp.current_process().name
    proclog=logging.getLogger(procName)
    proclog.setLevel(level=logging.INFO)
    mfh = QueueHandler(logqueue)
    mft = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    mfh.setFormatter(mft)
    proclog.addHandler(mfh)

    while work:
        #grabs project from queue
        try:
            projname = queue.get(block=True, timeout=3)
        except Queue.Empty:
            work=False
            proclog.info("exiting gracefully")
            break
        else:
            #locks the project : cannot be updated more than once.
            lockfile=os.path.join(options.lockdir, projname)
            if not os.path.exists(lockfile):
                try:
                    open(lockfile,'w').close()
                except:
                    proclog.error("cannot create lockfile {}".format(lockfile))
                try:
                    proj=mylims.get_projects(name=projname)[0]
                    P = PSUL(proj, samp_db, proj_db, options.upload, options.project_name, options.output_f, proclog)
                    P.handle_project()
                except :
                    error=sys.exc_info()
                    stack=traceback.extract_tb(error[2])
                    proclog.error("{0}:{1}\n{2}".format(error[0], error[1], formatStack(stack)))

                try:
                    os.remove(lockfile)
                except:
                    proclog.error("cannot remove lockfile {}".format(lockfile))

            #signals to queue job is done
            queue.task_done()

def masterProcess(options,projectList, mainlims, logger):
    projectsQueue=mp.JoinableQueue()
    logQueue=mp.Queue()
    childs=[]
    #Initial step : order projects by sample number:
    logger.info("ordering the project list")
    orderedprojectlist=sorted(projectList, key=lambda x: (mainlims.get_sample_number(projectname=x.name)), reverse=True)
    logger.info("done ordering the project list")
    #spawn a pool of processes, and pass them queue instance
    for i in range(options.processes):
        p = mp.Process(target=processPSUL, args=(options,projectsQueue, logQueue))
        p.start()
        childs.append(p)
    #populate queue with data
    for proj in orderedprojectlist:
        projectsQueue.put(proj.name)

    #wait on the queue until everything has been processed
    notDone=True
    while notDone:
        try:
            log=logQueue.get(False)
            logger.handle(log)
        except Queue.Empty:
            if not stillRunning(childs):
                notDone=False
                break

def stillRunning(processList):
    ret=False
    for p in processList:
        if p.is_alive():
            ret=True

    return ret

class QueueHandler(logging.Handler):
    """
    This handler sends events to a queue. Typically, it would be used together
    with a multiprocessing Queue to centralise logging to file in one process
    (in a multi-process application), so as to avoid file write contention
    between processes.

    This code is new in Python 3.2, but this class can be copy pasted into
    user code for use with earlier Python versions.
    """

    def __init__(self, queue):
        """
        Initialise an instance, using the passed queue.
        """
        logging.Handler.__init__(self)
        self.queue = queue

    def enqueue(self, record):
        """
        Enqueue a record.

        The base implementation uses put_nowait. You may want to override
        this method if you want to use blocking, timeouts or custom queue
        implementations.
        """
        self.queue.put_nowait(record)

    def prepare(self, record):
        """
        Prepares a record for queuing. The object returned by this method is
        enqueued.

        The base implementation formats the record to merge the message
        and arguments, and removes unpickleable items from the record
        in-place.

        You might want to override this method if you want to convert
        the record to a dict or JSON string, or send a modified copy
        of the record while leaving the original intact.
        """
        # The format operation gets traceback text into record.exc_text
        # (if there's exception data), and also puts the message into
        # record.message. We can then use this to replace the original
        # msg + args, as these might be unpickleable. We also zap the
        # exc_info attribute, as it's no longer needed and, if not None,
        # will typically not be pickleable.
        self.format(record)
        record.msg = record.message
        record.args = None
        record.exc_info = None
        return record

    def emit(self, record):
        """
        Emit a record.

        Writes the LogRecord to the queue, preparing it for pickling first.
        """
        try:
            self.enqueue(self.prepare(record))
        except Exception:
            self.handleError(record)

if __name__ == '__main__':
    usage = "Usage:       python project_summary_upload_LIMS.py [options]"
    parser = OptionParser(usage=usage)
    parser.add_option("-p", "--project", dest = "project_name", default = None,
                      help = "eg: M.Uhlen_13_01. Dont use with -a flagg.")
    parser.add_option("-a", "--all_projects", dest = "all_projects", action =
                      "store_true", default = False, help = ("Upload all Lims ",
                      "projects into couchDB. Don't use with -f flagg."))
    parser.add_option("-c", "--conf", dest = "conf", default = os.path.join(
                      os.environ['HOME'],'opt/config/post_process.yaml'), help =
                      "Config file.  Default: ~/opt/config/post_process.yaml")
    parser.add_option("--no_upload", dest = "upload", default = True, action =
                      "store_false", help = ("Use this tag if project objects ",
                      "should not be uploaded, but printed to output_f, or to ",
                      "stdout"))
    parser.add_option("--output_f", dest = "output_f", help = ("Output file",
                      " that will be used only if --no_upload tag is used"), default=None)
    parser.add_option("-m", "--multiprocs", type='int', dest = "processes", default = 4,
                      help = "How many processes will be spawned. Will only work with -a")
    parser.add_option("-l", "--logfile", dest = "logfile", help = ("log file",
                      " that will be used. default is $HOME/lims2db_projects.log "), default=os.path.expanduser("~/lims2db_projects.log"))
    parser.add_option("--lockdir", dest = "lockdir", help = ("directory handling the lock files",
                      " to avoid multiple updating of one project. default is $HOME/psul_locks "), default=os.path.expanduser("~/psul_locks"))
    parser.add_option("-j", "--hours", dest = "hours",type='int', help = ("only handle projects modified in the last X hours"), default=None)

    (options, args) = parser.parse_args()
    main(options)

<commit_msg>reworking update schema<commit_after>#!/usr/bin/env python
"""Script to load project info from Lims into the project database in statusdb.

Maya Brandi, Science for Life Laboratory, Stockholm, Sweden.
"""
from __future__ import print_function
from genologics.config import BASEURI, USERNAME, PASSWORD
from genologics.lims import *
from LIMS2DB.objectsDB.functions import *
from optparse import OptionParser
from statusdb.db.utils import *
from LIMS2DB.utils import formatStack

from pprint import pprint

import codecs
import datetime
import LIMS2DB.objectsDB.objectsDB as DB
import logging
import logging.handlers
import multiprocessing as mp
import os
import Queue
import sys
import time
import traceback

class PSUL():
    def __init__(self, proj, samp_db, proj_db, upload_data, man_name, output_f, log):
        self.proj = proj
        self.id = proj.id
        self.udfs = proj.udf
        self.name = proj.name
        self.open_date = proj.open_date
        self.close_date = proj.close_date
        self.samp_db = samp_db
        self.proj_db = proj_db
        self.upload_data = upload_data
        self.man_name = man_name
        self.output_f = output_f
        self.ordered_opened = self.get_ordered_opened()
        self.lims = Lims(BASEURI, USERNAME, PASSWORD)
        self.log=log

    def print_couchdb_obj_to_file(self, obj):
        if self.output_f is not None:
            with open(self.output_f, 'w') as f:
                print(obj, file = f)
        else:
            print(obj, file = sys.stdout)

    def get_ordered_opened(self):
        """Is project registered as opened or ordered?"""

        if self.open_date:
            return self.open_date
        elif 'Order received' in self.proj.udfs:
            return self.udfs['Order received'].isoformat()
        else:
            return datetime.date.today().strftime("%Y-%m-%d")

    def determine_update(self):
        """Determine wether to and how to update project"""
        opened_after_130630 = comp_dates('2013-06-30', self.ordered_opened)
        log_info = ''
        if not opened_after_130630 :
                start_update = False
                self.log.info('Project is not updated because: '
                    'It was opened or ordered before 2013-06-30 : '
                    '({ord_op})'.format(ord_op = self.ordered_opened))
        else:
            start_update = True

        return start_update

    def update_project(self):
        """Fetch project info and update project in the database."""
        self.log.info('Handling {proj}'.format(proj = self.name))
        project = DB.ProjectDB(self.lims, self.id, self.samp_db, self.log)

        key = find_proj_from_view(self.proj_db, self.name)
        project.obj['_id'] = find_or_make_key(key)
        if self.upload_data:
            info = save_couchdb_obj(self.proj_db, project.obj)
        else:
            info = self.print_couchdb_obj_to_file(project.obj)

        self.log.info("project {name} is handled and {info}: _id = {id}".format(
                           name=self.name, info=info, id=project.obj['_id']))

    def handle_project(self):
        if self.determine_update():
            self.update_project()

def main(options):
    conf = options.conf
    upload_data = options.upload
    output_f = options.output_f
    couch = load_couch_server(conf)
    proj_db = couch['projects']
    samp_db = couch['samples']
    mainlims = Lims(BASEURI, USERNAME, PASSWORD)
    mainlog = logging.getLogger('psullogger')
    mainlog.setLevel(level=logging.INFO)
    mfh = logging.handlers.RotatingFileHandler(options.logfile, maxBytes=209715200, backupCount=5)
    mft = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    mfh.setFormatter(mft)
    mainlog.addHandler(mfh)

    if options.all_projects:
        projects=create_projects_list(options, mainlims, mainlog)
        masterProcess(options,projects, mainlims, mainlog)
    elif options.project_name:
        proj = mainlims.get_projects(name = options.project_name)
        if not proj:
            mainlog.warn('No project named {man_name} in Lims'.format(
                        man_name = options.project_name))
        else:
            P = PSUL(proj[0], samp_db, proj_db, options.upload, options.project_name, output_f, mainlog)
            P.handle_project()

def create_projects_list(options, lims, log):
        projects = lims.get_projects()
        if options.hours:
            delta=datetime.timedelta(hours=options.hours)
            #timestring for processes is different from timestring for projects. First needs Z, second needs full timezone
            time_string_pc=(datetime.datetime.now()-delta).strftime('%Y-%m-%dT%H:%M:%SZ')
            time_string_pj=(datetime.datetime.now()-delta).strftime('%Y-%m-%dT%H:%M:%SCET')
            valid_projects=lims.get_projects(last_modified=time_string_pj)
            for proj in projects:
                procs=lims.get_processes(last_modified=time_string_pc, projectname=proj.name)
                if procs:
                    valid_projects.append(proj)
            log.info("project list : {0}".format(" ".join([p.id for p in valid_projects])))
            return valid_projects
        else:
            log.info("project list : {0}".format(" ".join([p.id for p in projects])))
            return projects

def processPSUL(options, queue, logqueue):
    couch = load_couch_server(options.conf)
    proj_db = couch['projects']
    samp_db = couch['samples']
    mylims = Lims(BASEURI, USERNAME, PASSWORD)
    work=True
    procName=mp.current_process().name
    proclog=logging.getLogger(procName)
    proclog.setLevel(level=logging.INFO)
    mfh = QueueHandler(logqueue)
    mft = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    mfh.setFormatter(mft)
    proclog.addHandler(mfh)

    while work:
        #grabs project from queue
        try:
            projname = queue.get(block=True, timeout=3)
            proclog.info("Approximately {} projects left in queue".format(queue.qsize()))
        except Queue.Empty:
            work=False
            proclog.info("exiting gracefully")
            break
        except NotImplementedError:
            #qsize failed, no big deal
            pass
        else:
            #locks the project : cannot be updated more than once.
            lockfile=os.path.join(options.lockdir, projname)
            if not os.path.exists(lockfile):
                try:
                    open(lockfile,'w').close()
                except:
                    proclog.error("cannot create lockfile {}".format(lockfile))
                try:
                    proj=mylims.get_projects(name=projname)[0]
                    P = PSUL(proj, samp_db, proj_db, options.upload, options.project_name, options.output_f, proclog)
                    P.handle_project()
                except :
                    error=sys.exc_info()
                    stack=traceback.extract_tb(error[2])
                    proclog.error("{0}:{1}\n{2}".format(error[0], error[1], formatStack(stack)))

                try:
                    os.remove(lockfile)
                except:
                    proclog.error("cannot remove lockfile {}".format(lockfile))

            #signals to queue job is done
            queue.task_done()

def masterProcess(options,projectList, mainlims, logger):
    projectsQueue=mp.JoinableQueue()
    logQueue=mp.Queue()
    childs=[]
    #Initial step : order projects by sample number:
    logger.info("ordering the project list")
    orderedprojectlist=sorted(projectList, key=lambda x: (mainlims.get_sample_number(projectname=x.name)), reverse=True)
    logger.info("done ordering the project list")
    #spawn a pool of processes, and pass them queue instance
    for i in range(options.processes):
        p = mp.Process(target=processPSUL, args=(options,projectsQueue, logQueue))
        p.start()
        childs.append(p)
    #populate queue with data
    for proj in orderedprojectlist:
        projectsQueue.put(proj.name)

    #wait on the queue until everything has been processed
    notDone=True
    while notDone:
        try:
            log=logQueue.get(False)
            logger.handle(log)
        except Queue.Empty:
            logger.info("Queue is empty, waiting on children")
            if not stillRunning(childs):
                notDone=False
                break

def stillRunning(processList):
    ret=False
    for p in processList:
        if p.is_alive():
            ret=True

    return ret

class QueueHandler(logging.Handler):
    """
    This handler sends events to a queue. Typically, it would be used together
    with a multiprocessing Queue to centralise logging to file in one process
    (in a multi-process application), so as to avoid file write contention
    between processes.

    This code is new in Python 3.2, but this class can be copy pasted into
    user code for use with earlier Python versions.
    """

    def __init__(self, queue):
        """
        Initialise an instance, using the passed queue.
        """
        logging.Handler.__init__(self)
        self.queue = queue

    def enqueue(self, record):
        """
        Enqueue a record.

        The base implementation uses put_nowait. You may want to override
        this method if you want to use blocking, timeouts or custom queue
        implementations.
        """
        self.queue.put_nowait(record)

    def prepare(self, record):
        """
        Prepares a record for queuing. The object returned by this method is
        enqueued.

        The base implementation formats the record to merge the message
        and arguments, and removes unpickleable items from the record
        in-place.

        You might want to override this method if you want to convert
        the record to a dict or JSON string, or send a modified copy
        of the record while leaving the original intact.
        """
        # The format operation gets traceback text into record.exc_text
        # (if there's exception data), and also puts the message into
        # record.message. We can then use this to replace the original
        # msg + args, as these might be unpickleable. We also zap the
        # exc_info attribute, as it's no longer needed and, if not None,
        # will typically not be pickleable.
        self.format(record)
        record.msg = record.message
        record.args = None
        record.exc_info = None
        return record

    def emit(self, record):
        """
        Emit a record.

        Writes the LogRecord to the queue, preparing it for pickling first.
        """
        try:
            self.enqueue(self.prepare(record))
        except Exception:
            self.handleError(record)

if __name__ == '__main__':
    usage = "Usage:       python project_summary_upload_LIMS.py [options]"
    parser = OptionParser(usage=usage)
    parser.add_option("-p", "--project", dest = "project_name", default = None,
                      help = "eg: M.Uhlen_13_01. Dont use with -a flagg.")
    parser.add_option("-a", "--all_projects", dest = "all_projects", action =
                      "store_true", default = False, help = ("Upload all Lims ",
                      "projects into couchDB. Don't use with -f flagg."))
    parser.add_option("-c", "--conf", dest = "conf", default = os.path.join(
                      os.environ['HOME'],'opt/config/post_process.yaml'), help =
                      "Config file.  Default: ~/opt/config/post_process.yaml")
    parser.add_option("--no_upload", dest = "upload", default = True, action =
                      "store_false", help = ("Use this tag if project objects ",
                      "should not be uploaded, but printed to output_f, or to ",
                      "stdout"))
    parser.add_option("--output_f", dest = "output_f", help = ("Output file",
                      " that will be used only if --no_upload tag is used"), default=None)
    parser.add_option("-m", "--multiprocs", type='int', dest = "processes", default = 4,
                      help = "How many processes will be spawned. Will only work with -a")
    parser.add_option("-l", "--logfile", dest = "logfile", help = ("log file",
                      " that will be used. default is $HOME/lims2db_projects.log "), default=os.path.expanduser("~/lims2db_projects.log"))
    parser.add_option("--lockdir", dest = "lockdir", help = ("directory handling the lock files",
                      " to avoid multiple updating of one project. default is $HOME/psul_locks "), default=os.path.expanduser("~/psul_locks"))
    parser.add_option("-j", "--hours", dest = "hours",type='int', help = ("only handle projects modified in the last X hours"), default=None)
    parser.add_option("-k", "--control", dest = "control", action="store_true", help = ("only perform a dry-run"), default=False)

    (options, args) = parser.parse_args()
    main(options)

<|endoftext|>


<commit_before>"""Operations regarding miscellaneous files."""

from __future__ import absolute_import, division, print_function
import os.path
import re

import numpy as np

COMMENT_PREFIX = '#'

def toline(iterable):
    """Convert an iterable into a line."""
    return ' '.join(str(x) for x in iterable) + '\n'

def valid_lines(filename):
    """Read and yield lines that are neither empty nor comments."""
    with open(filename, 'rU') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith(COMMENT_PREFIX):
                yield line

def read_subwindows(filename):
    """Read a list of subwindows from file, return in a dictionary."""
    r = {}
    for line in valid_lines(filename):
        words = line.split()
        if len(words) != 8:
            raise Exception('Cannot parse subwindow: %s' % line)
        case, scan = int(words[0]), words[1]
        subwindow = [int(x) for x in words[2:]]
        r[(case, scan)] = subwindow
    return r

def read_patients_file(filename, include_lines=False):
    """Load a list of patients.

    Row format: num name scan1,scan2,... score [location]
    """
    from dwi.patient import GleasonScore, Lesion, Patient
    patients = []
    p = re.compile(r"""
                   (?P<num>\d+) \s+
                   (?P<name>\w+) \s+
                   (?P<scans>[\w,]+) \s+
                   (?P<score>\d\+\d(\+\d)?) \s* (?P<location>\w+)? \s*
                   ((?P<score2>\d\+\d(\+\d)?) \s+ (?P<location2>\w+))? \s*
                   ((?P<score3>\d\+\d(\+\d)?) \s+ (?P<location3>\w+))?
                   """,
                   flags=re.VERBOSE)
    for line in valid_lines(filename):
        m = p.match(line)
        if m:
            num = int(m.group('num'))
            name = m.group('name').lower()
            scans = sorted(m.group('scans').lower().split(','))
            score = GleasonScore(m.group('score'))
            lesions = [Lesion(0, score, 'xx')]
            if m.group('location'):
                # New-style, multi-lesion file.
                lesions = []
                lesions.append(Lesion(0, GleasonScore(m.group('score')),
                                      m.group('location').lower()))
                if m.group('score2'):
                    lesions.append(Lesion(1, GleasonScore(m.group('score2')),
                                          m.group('location2').lower()))
                if m.group('score3'):
                    lesions.append(Lesion(2, GleasonScore(m.group('score3')),
                                          m.group('location3').lower()))
            patient = Patient(num, name, scans, lesions)
            if include_lines:
                patient.line = line
            patients.append(patient)
        else:
            raise Exception('Invalid line in patients file: %s', line)
    return sorted(patients)

def read_sample_list(filename):
    """Read a list of samples from file."""
    entries = []
    p = re.compile(r'(\d+)\s+(\w+)\s+([\w,]+)')
    for line in valid_lines(filename):
        m = p.match(line)
        if m:
            case, name, scans = m.groups()
            case = int(case)
            name = name.lower()
            scans = tuple(sorted(scans.lower().split(',')))
            d = dict(case=case, name=name, scans=scans)
            entries.append(d)
    return entries

def read_mask_file(filename):
    """Read a ROI mask file. XXX: Deprecated, use dwi.mask module instead."""
    arrays = []
    for line in valid_lines(filename):
        if line[0] == '0' or line[0] == '1':
            a = np.array(list(line), dtype=int)
            arrays.append(a)
    mask = np.array(arrays)
    return mask

def read_subregion_file(filename):
    """Read a subregion definition from file.

    It's formatted as one voxel index per line, zero-based, in order of
    y_first, y_last, x_first, x_last, z_first, z_last. The "last" ones need +1
    to get Python-like start:stop indices. They are returned as (z_start,
    z_stop, y_start, y_stop, x_start, x_stop).
    """
    entries = []
    for line in valid_lines(filename):
        entries.append(int(float(line)))
    if len(entries) != 6:
        raise Exception('Invalid subregion file: %s' % filename)
    entries = entries[4:6] + entries[0:4]  # Move z indices to front.
    # Add one to "last" indices get Python-like start:stop indices.
    entries[1] += 1
    entries[3] += 1
    entries[5] += 1
    return tuple(entries)

def write_comment(f, text):
    """Write zero or more lines to file with comment prefix."""
    for line in text.splitlines():
        f.write('{p} {s}\n'.format(p=COMMENT_PREFIX, s=line))

def write_subregion_file(filename, win, comment=''):
    """Write a subregion definition to file.

    It's formatted as one voxel index per line, zero-based, in order of
    y_first, y_last, x_first, x_last, z_first, z_last.
    """
    if len(win) != 6:
        raise Exception('Invalid subregion: %s' % win)
    entries = [win[2], win[3]-1, win[4], win[5]-1, win[0], win[1]-1]
    with open(filename, 'w') as f:
        write_comment(f, comment)
        for entry in entries:
            f.write('%i\n' % entry)

def write_pmap(filename, pmap, attrs, fmt=None):
    """Write parametric map file either as HDF5 or ASCII."""
    pmap = np.asanyarray(pmap)
    if pmap.ndim < 2:
        raise Exception('Not enough dimensions: {}'.format(pmap.shape))
    if 'parameters' not in attrs:
        attrs['parameters'] = [str(i) for i in range(pmap.shape[-1])]
    if 'shape' not in attrs:
        attrs['shape'] = pmap.shape
    if 'dtype' not in attrs:
        attrs['dtype'] = str(pmap.dtype)
    if pmap.shape[-1] != len(attrs['parameters']):
        raise Exception('Number of values and parameters mismatch')
    assert all(isinstance(x, str) for x in
               attrs['parameters']), attrs['parameters']
    if fmt is None:
        fmt = os.path.splitext(filename)[1][1:]
    if fmt in ['hdf5', 'h5']:
        import dwi.hdf5
        dwi.hdf5.write_hdf5(filename, pmap, attrs)
    elif fmt in ['txt', 'ascii']:
        import dwi.asciifile
        pmap = pmap.reshape((-1, pmap.shape[-1]))  # Can't keep shape.
        dwi.asciifile.write_ascii_file(filename, pmap, None, attrs=attrs)
    else:
        raise Exception('Unknown format: {}'.format(fmt))

def pick_params(pmap, attrs, params):
    """Select a subset of parameters by their indices."""
    params = list(params)
    pmap = pmap[..., params]
    if len(attrs['bset']) == len(attrs['parameters']):
        attrs['bset'] = [attrs['bset'][x] for x in params]
    if len(attrs['echotimes']) == len(attrs['parameters']):
        attrs['echotimes'] = [attrs['echotimes'][x] for x in params]
    attrs['parameters'] = [attrs['parameters'][x] for x in params]
    return pmap, attrs

def read_pmap(pathname, ondisk=False, fmt=None, params=None):
    """Read a parametric map.

    With parameter ondisk it will not be read into memory. Parameter params
    tells which parameter indices should be included.
    """
    if fmt is None:
        fmt = os.path.splitext(pathname)[1][1:]
    if fmt in ['hdf5', 'h5']:
        import dwi.hdf5
        pmap, attrs = dwi.hdf5.read_hdf5(pathname, ondisk=ondisk)
    elif fmt in ['txt', 'ascii']:
        import dwi.asciifile
        attrs, pmap = dwi.asciifile.read_ascii_file(pathname)
        if 'parameters' in attrs:
            attrs['parameters'] = attrs['parameters'].split()
    else:
        import dwi.dicomfile
        d = dwi.dicomfile.read_dir(pathname)
        pmap = d.pop('image')
        attrs = dict(d)
    if 'parameters' not in attrs:
        attrs['parameters'] = range(pmap.shape[-1])
    attrs['parameters'] = [str(x) for x in attrs['parameters']]
    if params is not None:
        pmap, attrs = pick_params(pmap, attrs, params)
    return pmap, attrs
<commit_msg>handle missing attributes<commit_after>"""Operations regarding miscellaneous files."""

from __future__ import absolute_import, division, print_function
import os.path
import re

import numpy as np

COMMENT_PREFIX = '#'

def toline(iterable):
    """Convert an iterable into a line."""
    return ' '.join(str(x) for x in iterable) + '\n'

def valid_lines(filename):
    """Read and yield lines that are neither empty nor comments."""
    with open(filename, 'rU') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith(COMMENT_PREFIX):
                yield line

def read_subwindows(filename):
    """Read a list of subwindows from file, return in a dictionary."""
    r = {}
    for line in valid_lines(filename):
        words = line.split()
        if len(words) != 8:
            raise Exception('Cannot parse subwindow: %s' % line)
        case, scan = int(words[0]), words[1]
        subwindow = [int(x) for x in words[2:]]
        r[(case, scan)] = subwindow
    return r

def read_patients_file(filename, include_lines=False):
    """Load a list of patients.

    Row format: num name scan1,scan2,... score [location]
    """
    from dwi.patient import GleasonScore, Lesion, Patient
    patients = []
    p = re.compile(r"""
                   (?P<num>\d+) \s+
                   (?P<name>\w+) \s+
                   (?P<scans>[\w,]+) \s+
                   (?P<score>\d\+\d(\+\d)?) \s* (?P<location>\w+)? \s*
                   ((?P<score2>\d\+\d(\+\d)?) \s+ (?P<location2>\w+))? \s*
                   ((?P<score3>\d\+\d(\+\d)?) \s+ (?P<location3>\w+))?
                   """,
                   flags=re.VERBOSE)
    for line in valid_lines(filename):
        m = p.match(line)
        if m:
            num = int(m.group('num'))
            name = m.group('name').lower()
            scans = sorted(m.group('scans').lower().split(','))
            score = GleasonScore(m.group('score'))
            lesions = [Lesion(0, score, 'xx')]
            if m.group('location'):
                # New-style, multi-lesion file.
                lesions = []
                lesions.append(Lesion(0, GleasonScore(m.group('score')),
                                      m.group('location').lower()))
                if m.group('score2'):
                    lesions.append(Lesion(1, GleasonScore(m.group('score2')),
                                          m.group('location2').lower()))
                if m.group('score3'):
                    lesions.append(Lesion(2, GleasonScore(m.group('score3')),
                                          m.group('location3').lower()))
            patient = Patient(num, name, scans, lesions)
            if include_lines:
                patient.line = line
            patients.append(patient)
        else:
            raise Exception('Invalid line in patients file: %s', line)
    return sorted(patients)

def read_sample_list(filename):
    """Read a list of samples from file."""
    entries = []
    p = re.compile(r'(\d+)\s+(\w+)\s+([\w,]+)')
    for line in valid_lines(filename):
        m = p.match(line)
        if m:
            case, name, scans = m.groups()
            case = int(case)
            name = name.lower()
            scans = tuple(sorted(scans.lower().split(',')))
            d = dict(case=case, name=name, scans=scans)
            entries.append(d)
    return entries

def read_mask_file(filename):
    """Read a ROI mask file. XXX: Deprecated, use dwi.mask module instead."""
    arrays = []
    for line in valid_lines(filename):
        if line[0] == '0' or line[0] == '1':
            a = np.array(list(line), dtype=int)
            arrays.append(a)
    mask = np.array(arrays)
    return mask

def read_subregion_file(filename):
    """Read a subregion definition from file.

    It's formatted as one voxel index per line, zero-based, in order of
    y_first, y_last, x_first, x_last, z_first, z_last. The "last" ones need +1
    to get Python-like start:stop indices. They are returned as (z_start,
    z_stop, y_start, y_stop, x_start, x_stop).
    """
    entries = []
    for line in valid_lines(filename):
        entries.append(int(float(line)))
    if len(entries) != 6:
        raise Exception('Invalid subregion file: %s' % filename)
    entries = entries[4:6] + entries[0:4]  # Move z indices to front.
    # Add one to "last" indices get Python-like start:stop indices.
    entries[1] += 1
    entries[3] += 1
    entries[5] += 1
    return tuple(entries)

def write_comment(f, text):
    """Write zero or more lines to file with comment prefix."""
    for line in text.splitlines():
        f.write('{p} {s}\n'.format(p=COMMENT_PREFIX, s=line))

def write_subregion_file(filename, win, comment=''):
    """Write a subregion definition to file.

    It's formatted as one voxel index per line, zero-based, in order of
    y_first, y_last, x_first, x_last, z_first, z_last.
    """
    if len(win) != 6:
        raise Exception('Invalid subregion: %s' % win)
    entries = [win[2], win[3]-1, win[4], win[5]-1, win[0], win[1]-1]
    with open(filename, 'w') as f:
        write_comment(f, comment)
        for entry in entries:
            f.write('%i\n' % entry)

def write_pmap(filename, pmap, attrs, fmt=None):
    """Write parametric map file either as HDF5 or ASCII."""
    pmap = np.asanyarray(pmap)
    if pmap.ndim < 2:
        raise Exception('Not enough dimensions: {}'.format(pmap.shape))
    if 'parameters' not in attrs:
        attrs['parameters'] = [str(i) for i in range(pmap.shape[-1])]
    if 'shape' not in attrs:
        attrs['shape'] = pmap.shape
    if 'dtype' not in attrs:
        attrs['dtype'] = str(pmap.dtype)
    if pmap.shape[-1] != len(attrs['parameters']):
        raise Exception('Number of values and parameters mismatch')
    assert all(isinstance(x, str) for x in
               attrs['parameters']), attrs['parameters']
    if fmt is None:
        fmt = os.path.splitext(filename)[1][1:]
    if fmt in ['hdf5', 'h5']:
        import dwi.hdf5
        dwi.hdf5.write_hdf5(filename, pmap, attrs)
    elif fmt in ['txt', 'ascii']:
        import dwi.asciifile
        pmap = pmap.reshape((-1, pmap.shape[-1]))  # Can't keep shape.
        dwi.asciifile.write_ascii_file(filename, pmap, None, attrs=attrs)
    else:
        raise Exception('Unknown format: {}'.format(fmt))

def pick_params(pmap, attrs, params):
    """Select a subset of parameters by their indices."""
    params = list(params)
    pmap = pmap[..., params]
    if 'bset' in attrs and len(attrs['bset']) == len(attrs['parameters']):
        attrs['bset'] = [attrs['bset'][x] for x in params]
    if 'echotimes' in attrs and (len(attrs['echotimes']) ==
                                 len(attrs['parameters'])):
        attrs['echotimes'] = [attrs['echotimes'][x] for x in params]
    attrs['parameters'] = [attrs['parameters'][x] for x in params]
    return pmap, attrs

def read_pmap(pathname, ondisk=False, fmt=None, params=None):
    """Read a parametric map.

    With parameter ondisk it will not be read into memory. Parameter params
    tells which parameter indices should be included.
    """
    if fmt is None:
        fmt = os.path.splitext(pathname)[1][1:]
    if fmt in ['hdf5', 'h5']:
        import dwi.hdf5
        pmap, attrs = dwi.hdf5.read_hdf5(pathname, ondisk=ondisk)
    elif fmt in ['txt', 'ascii']:
        import dwi.asciifile
        attrs, pmap = dwi.asciifile.read_ascii_file(pathname)
        if 'parameters' in attrs:
            attrs['parameters'] = attrs['parameters'].split()
    else:
        import dwi.dicomfile
        d = dwi.dicomfile.read_dir(pathname)
        pmap = d.pop('image')
        attrs = dict(d)
    if 'parameters' not in attrs:
        attrs['parameters'] = range(pmap.shape[-1])
    attrs['parameters'] = [str(x) for x in attrs['parameters']]
    if params is not None:
        pmap, attrs = pick_params(pmap, attrs, params)
    return pmap, attrs
<|endoftext|>


<commit_before>'''Fit an L1-Regularized Reversible Gaussian Hidden Markov Model
'''
# Author: Robert McGibbon <rmcgibbo@gmail.com>
# Contributors:
# Copyright (c) 2014, Stanford University
# All rights reserved.

# Mixtape is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with Mixtape. If not, see <http://www.gnu.org/licenses/>.

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

from __future__ import print_function, division, absolute_import
import sys
import os
import glob
import json
import time
import numpy as np
import mdtraj as md

# from sklearn.cross_validation import KFold
from mixtape.ghmm import GaussianFusionHMM
# from mixtape.lagtime import contraction
from mixtape.cmdline import Command, argument_group, MultipleIntAction
from mixtape.commands.mixins import MDTrajInputMixin
import mixtape.featurizer

__all__ = ['FitGHMM']

#-----------------------------------------------------------------------------
# Code
#-----------------------------------------------------------------------------

class FitGHMM(Command, MDTrajInputMixin):
    name = 'fit-ghmm'
    description = '''Fit L1-Regularized Reversible Gaussian hidden Markov models with EM.'''

    group_hmm = argument_group('HMM Options')
    group_hmm.add_argument('--featurizer', type=str, required=True, help='''
        Path to saved featurizer object''')
    group_hmm.add_argument('-k', '--n-states', action=MultipleIntAction, default=[2], help='''
        Number of states in the models. Default = [2,]''', nargs='+')
    group_hmm.add_argument('-l', '--lag-times', action=MultipleIntAction, default=[1], help='''
        Lag time(s) of the model(s). Default = [1,]''', nargs='+')
    group_hmm.add_argument('--platform', choices=['cuda', 'cpu', 'sklearn'],
        default='cpu', help='Implementation platform. default="cpu"')
    group_hmm.add_argument('--fusion-prior', type=float, default=1e-2, help='''
        Strength of the adaptive fusion prior. default=1e-2''')
    group_hmm.add_argument('--n-init', type=int, default=10, help='''Number
        of initialization for each model fit. Each of these "outer iterations"
        corresponds to a new random initialization of the states from kmeans
        and then `--n-em-iter`s of expectation-maximization. The best of these
        models (selected by likelihood) is retained. default=10''')
    group_hmm.add_argument('--n-em-iter', type=int, default=10, help='''
        Maximum number of iterations of expectation-maximization steps
        per fitting round. default=10''')
    group_hmm.add_argument('--thresh', type=float, default=1e-2, help='''
        Convergence criterion for EM. Quit when the log likelihood
        decreases by less than this threshold. default=1e-2''')
    group_hmm.add_argument('--n-lqa-iter', type=int, default=10, help='''
        Max number of iterations for local quadradric approximation
        solving the fusion-L1. default=10''')
    group_hmm.add_argument('--reversible-type', choices=['mle', 'transpose'], help='''
        Method by which the model is constrained to be
        reversible. default="mle"''', default='mle')
    group_hmm.add_argument('-sp', '--split', type=int, help='''Split
        trajectories into smaller chunks. This looses some counts (i.e. like
        1%% of the counts are lost with --split 100), but can help with speed
        (on gpu + multicore cpu) and numerical instabilities that come when
        trajectories get extremely long.''', default=10000)
    group_hmm.add_argument('-n-reps', '--n-repetitions', type=int, default=1, help='''
        Run this many repetitions of the *entire experiment*  with
        different random seeds, saving all of the models to the output file.
        This is the outer-most iteration. default=1''')

    # group_cv = argument_group('Cross Validation')
    # group_cv.add_argument('--n-cv', type=int, default=1,
    #     help='Run N-fold cross validation. default=1')
    # We're training and testing at the same lag time for the moment
    # group_cv.add_argument('--test-lag-time', type=int, default=1,
    #     help='Lag time at which to test the models. default=1')

    group_out = argument_group('Output')
    group_out.add_argument('-o', '--out', default='hmms.jsonlines',
                           help='Output file. default="hmms.jsonlines"')

    def __init__(self, args):
        self.args = args
        if args.top is not None:
            self.top = md.load(os.path.expanduser(args.top))
        else:
            self.top = None

        self.featurizer = mixtape.featurizer.load(args.featurizer)
        self.filenames = glob.glob(os.path.expanduser(args.dir) + '/*.' + args.ext)
        self.n_features = self.featurizer.n_features

    def start(self):
        args = self.args
        data = self.load_data()

        with open(args.out, 'a', 1) as outfile:
            outfile.write('# %s\n' % ' '.join(sys.argv))

            for lag_time in args.lag_times:
                subsampled = [d[::lag_time] for d in data]
                for n_states in args.n_states:
                    for rep in range(self.args.n_repetitions):
                        self.fit(subsampled, subsampled, n_states, lag_time, rep, args, outfile)

                    # if args.n_cv > 1:
                    # for fold, (train_i, test_i) in enumerate(KFold(n=len(data), n_folds=args.n_cv)):
                    #        train = [subsampled[i] for i in train_i]
                    #        test = [subsampled[i] for i in test_i]
                    #
                    #        self.fit(train, test, n_states, lag_time, fold, args, outfile)
                    # else:
                    #     self.fit(subsampled, subsampled, n_states, lag_time, 0, args, outfile)

    def fit(self, train, test, n_states, train_lag_time, repetition, args, outfile):
        kwargs = dict(n_states=n_states, n_features=self.n_features,
                      n_init=args.n_init,
                      n_em_iter=args.n_em_iter, n_lqa_iter=args.n_lqa_iter,
                      fusion_prior=args.fusion_prior, thresh=args.thresh,
                      reversible_type=args.reversible_type, platform=args.platform)
        model = GaussianFusionHMM(**kwargs)

        start = time.time()
        model.fit(train)
        end = time.time()

        result = {
            'model': 'GaussianFusionHMM',
            'timescales': (np.real(model.timescales_) * train_lag_time).tolist(),
            'transmat': np.real(model.transmat_).tolist(),
            'populations': np.real(model.populations_).tolist(),
            'n_states': model.n_states,
            'split': args.split,
            'fusion_prior': args.fusion_prior,
            'train_lag_time': train_lag_time,
            'train_time': end - start,
            'means': np.real(model.means_).tolist(),
            'vars': np.real(model.vars_).tolist(),
            'train_logprob': model.fit_logprob_[-1],
            'n_train_observations': sum(len(t) for t in train),
            'n_test_observations': sum(len(t) for t in test),
            'train_logprobs': model.fit_logprob_,
            #'test_lag_time': args.test_lag_time,
            'cross_validation_fold': 0,
            'cross_validation_nfolds': 1,
            'repetition': repetition,
        }

        # model.transmat_ = contraction(model.transmat_, float(train_lag_time) / float(args.test_lag_time))
        # Don't do any contraction -- train and test at the same lagtime
        result['test_logprob'] = model.score(test)
        result['test_lag_time'] = train_lag_time

        if not np.all(np.isfinite(model.transmat_)):
            print('Nonfinite numbers in transmat !!')

        json.dump(result, outfile)
        outfile.write('\n')

    def load_data(self):
        load_time_start = time.time()
        data = []
        for tfn in self.filenames:
            kwargs = {} if tfn.endswith('h5') else {'top': self.top}
            for t in md.iterload(tfn, chunk=self.args.split, **kwargs):
                features = self.featurizer.partial_transform(t)
                data.append(features)

        print('Loading data into memory + vectorization: %f s' % (time.time() - load_time_start))
        print('Fitting with %s timeseries from %d trajectories with %d total observations' % (
            len(data), len(self.filenames), sum(len(e) for e in data)))

        return data
<commit_msg>fitghmm doesnt need to reference featurizer.n_features<commit_after>'''Fit an L1-Regularized Reversible Gaussian Hidden Markov Model
'''
# Author: Robert McGibbon <rmcgibbo@gmail.com>
# Contributors:
# Copyright (c) 2014, Stanford University
# All rights reserved.

# Mixtape is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with Mixtape. If not, see <http://www.gnu.org/licenses/>.

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

from __future__ import print_function, division, absolute_import
import sys
import os
import glob
import json
import time
import numpy as np
import mdtraj as md

# from sklearn.cross_validation import KFold
from mixtape.ghmm import GaussianFusionHMM
# from mixtape.lagtime import contraction
from mixtape.cmdline import Command, argument_group, MultipleIntAction
from mixtape.commands.mixins import MDTrajInputMixin
import mixtape.featurizer

__all__ = ['FitGHMM']

#-----------------------------------------------------------------------------
# Code
#-----------------------------------------------------------------------------

class FitGHMM(Command, MDTrajInputMixin):
    name = 'fit-ghmm'
    description = '''Fit L1-Regularized Reversible Gaussian hidden Markov models with EM.'''

    group_hmm = argument_group('HMM Options')
    group_hmm.add_argument('--featurizer', type=str, required=True, help='''
        Path to saved featurizer object''')
    group_hmm.add_argument('-k', '--n-states', action=MultipleIntAction, default=[2], help='''
        Number of states in the models. Default = [2,]''', nargs='+')
    group_hmm.add_argument('-l', '--lag-times', action=MultipleIntAction, default=[1], help='''
        Lag time(s) of the model(s). Default = [1,]''', nargs='+')
    group_hmm.add_argument('--platform', choices=['cuda', 'cpu', 'sklearn'],
        default='cpu', help='Implementation platform. default="cpu"')
    group_hmm.add_argument('--fusion-prior', type=float, default=1e-2, help='''
        Strength of the adaptive fusion prior. default=1e-2''')
    group_hmm.add_argument('--n-init', type=int, default=10, help='''Number
        of initialization for each model fit. Each of these "outer iterations"
        corresponds to a new random initialization of the states from kmeans
        and then `--n-em-iter`s of expectation-maximization. The best of these
        models (selected by likelihood) is retained. default=10''')
    group_hmm.add_argument('--n-em-iter', type=int, default=10, help='''
        Maximum number of iterations of expectation-maximization steps
        per fitting round. default=10''')
    group_hmm.add_argument('--thresh', type=float, default=1e-2, help='''
        Convergence criterion for EM. Quit when the log likelihood
        decreases by less than this threshold. default=1e-2''')
    group_hmm.add_argument('--n-lqa-iter', type=int, default=10, help='''
        Max number of iterations for local quadradric approximation
        solving the fusion-L1. default=10''')
    group_hmm.add_argument('--reversible-type', choices=['mle', 'transpose'], help='''
        Method by which the model is constrained to be
        reversible. default="mle"''', default='mle')
    group_hmm.add_argument('-sp', '--split', type=int, help='''Split
        trajectories into smaller chunks. This looses some counts (i.e. like
        1%% of the counts are lost with --split 100), but can help with speed
        (on gpu + multicore cpu) and numerical instabilities that come when
        trajectories get extremely long.''', default=10000)
    group_hmm.add_argument('-n-reps', '--n-repetitions', type=int, default=1, help='''
        Run this many repetitions of the *entire experiment*  with
        different random seeds, saving all of the models to the output file.
        This is the outer-most iteration. default=1''')

    # group_cv = argument_group('Cross Validation')
    # group_cv.add_argument('--n-cv', type=int, default=1,
    #     help='Run N-fold cross validation. default=1')
    # We're training and testing at the same lag time for the moment
    # group_cv.add_argument('--test-lag-time', type=int, default=1,
    #     help='Lag time at which to test the models. default=1')

    group_out = argument_group('Output')
    group_out.add_argument('-o', '--out', default='hmms.jsonlines',
                           help='Output file. default="hmms.jsonlines"')

    def __init__(self, args):
        self.args = args
        if args.top is not None:
            self.top = md.load(os.path.expanduser(args.top))
        else:
            self.top = None

        self.featurizer = mixtape.featurizer.load(args.featurizer)
        self.filenames = glob.glob(os.path.expanduser(args.dir) + '/*.' + args.ext)

    def start(self):
        args = self.args
        data = self.load_data()

        with open(args.out, 'a', 1) as outfile:
            outfile.write('# %s\n' % ' '.join(sys.argv))

            for lag_time in args.lag_times:
                subsampled = [d[::lag_time] for d in data]
                for n_states in args.n_states:
                    for rep in range(self.args.n_repetitions):
                        self.fit(subsampled, subsampled, n_states, lag_time, rep, args, outfile)

                    # if args.n_cv > 1:
                    # for fold, (train_i, test_i) in enumerate(KFold(n=len(data), n_folds=args.n_cv)):
                    #        train = [subsampled[i] for i in train_i]
                    #        test = [subsampled[i] for i in test_i]
                    #
                    #        self.fit(train, test, n_states, lag_time, fold, args, outfile)
                    # else:
                    #     self.fit(subsampled, subsampled, n_states, lag_time, 0, args, outfile)

    def fit(self, train, test, n_states, train_lag_time, repetition, args, outfile):
        n_features = train[0].shape[1]
        kwargs = dict(n_states=n_states, n_features=n_features,
                      n_init=args.n_init,
                      n_em_iter=args.n_em_iter, n_lqa_iter=args.n_lqa_iter,
                      fusion_prior=args.fusion_prior, thresh=args.thresh,
                      reversible_type=args.reversible_type, platform=args.platform)
        model = GaussianFusionHMM(**kwargs)

        start = time.time()
        model.fit(train)
        end = time.time()

        result = {
            'model': 'GaussianFusionHMM',
            'timescales': (np.real(model.timescales_) * train_lag_time).tolist(),
            'transmat': np.real(model.transmat_).tolist(),
            'populations': np.real(model.populations_).tolist(),
            'n_states': model.n_states,
            'split': args.split,
            'fusion_prior': args.fusion_prior,
            'train_lag_time': train_lag_time,
            'train_time': end - start,
            'means': np.real(model.means_).tolist(),
            'vars': np.real(model.vars_).tolist(),
            'train_logprob': model.fit_logprob_[-1],
            'n_train_observations': sum(len(t) for t in train),
            'n_test_observations': sum(len(t) for t in test),
            'train_logprobs': model.fit_logprob_,
            #'test_lag_time': args.test_lag_time,
            'cross_validation_fold': 0,
            'cross_validation_nfolds': 1,
            'repetition': repetition,
        }

        # model.transmat_ = contraction(model.transmat_, float(train_lag_time) / float(args.test_lag_time))
        # Don't do any contraction -- train and test at the same lagtime
        result['test_logprob'] = model.score(test)
        result['test_lag_time'] = train_lag_time

        if not np.all(np.isfinite(model.transmat_)):
            print('Nonfinite numbers in transmat !!')

        json.dump(result, outfile)
        outfile.write('\n')

    def load_data(self):
        load_time_start = time.time()
        data = []
        for tfn in self.filenames:
            kwargs = {} if tfn.endswith('h5') else {'top': self.top}
            for t in md.iterload(tfn, chunk=self.args.split, **kwargs):
                features = self.featurizer.partial_transform(t)
                data.append(features)

        print('Loading data into memory + vectorization: %f s' % (time.time() - load_time_start))
        print('Fitting with %s timeseries from %d trajectories with %d total observations' % (
            len(data), len(self.filenames), sum(len(e) for e in data)))

        return data
<|endoftext|>