<commit_before>"""Definition of ephemeris
"""

import numpy as np
from datetime import timedelta

from .listeners import Speaker
from ..frames.frames import orbit2frame

class Ephem(Speaker):
    """This class represents a range of orbits

    It provides several useful interfaces in order to compute through them

    Example:

    .. code-block:: python

        ephem = orb.ephem(Date.now(), timedelta(hours=6), timedelta(minutes=2))
        ephem.frame = 'ITRF'
        ephem.form = 'spherical'
        latitudes = ephem[:,1]
        longitudes = ephem[:,2]
    """

    LINEAR = "linear"
    LAGRANGE = "lagrange"
    DEFAULT_ORDER = 8

    def __init__(self, orbits, method=None, order=None):
        self._orbits = list(sorted(orbits, key=lambda x: x.date))
        self.method = self.LAGRANGE if method is None else method
        self.order = order if isinstance(order, int) else self.DEFAULT_ORDER

    def __iter__(self):
        self._i = -1
        return self

    def __next__(self):
        self._i += 1
        if self._i >= len(self._orbits):
            raise StopIteration
        return self._orbits[self._i]

    def __getitem__(self, index):
        if isinstance(index, (slice, int)):
            # retrieve a defined Orbit object
            return self._orbits[index]
        else:
            # In order to comply with numpy array handling (e.g. column selection)
            return np.array(self._orbits)[index]

    def __len__(self):
        return len(self._orbits)

    @property
    def start(self):
        """Date of the first element
        """
        return self._orbits[0].date

    @property
    def stop(self):
        """Date of the last element
        """
        return self._orbits[-1].date

    @property
    def dates(self):
        """Generator yielding Dates of each Orbit object of the ephem
        """
        return (o.date for o in self)

    # @property
    # def steps(self):
    #     """Time intervals used in the ephemeris
    #     """
    #     return set([self[i].date - self[i + 1].date for i in range(len(self) - 1)])

    @property
    def frame(self):
        """Get the frame of the first point
        """
        return self._orbits[0].frame

    @frame.setter
    def frame(self, frame):  # pragma: no cover
        """Change the frames of all points
        """
        for orb in self:
            orb.frame = frame

    @property
    def form(self):  # pragma: no cover
        """Get the form of the first point
        """
        return self._orbits[0].form

    @form.setter
    def form(self, form):  # pragma: no cover
        """Change the form of all points
        """
        for orb in self:
            orb.form = form

    def interpolate(self, date, method=None, order=None):
        """Interpolate data at a given date

        Args:
            date (Date):
            method (str): Method of interpolation to use
            order (int): In case of ``LAGRANGE`` method is used
        Return:
            Orbit:
        Raise:
            ValueError: when date is not in the range of the ephemeris
            ValueError: when the order of interpolation is insufficient
        """

        if not self.start <= date <= self.stop:
            raise ValueError("Date '%s' not in range" % date)

        prev_idx = 0
        ephem = self

        # Binary search of the orbit step just before the desired date
        while True:
            idx = len(ephem)
            if idx == 1:
                break
            k = idx // 2

            if date > ephem[k].date:
                prev_idx += k
                ephem = ephem[k:]
            else:
                ephem = ephem[:k]

        method = method if method is not None else self.method
        order = order if order is not None else self.order

        if method == self.LINEAR:

            y0 = self[prev_idx]
            y1 = self[prev_idx + 1]

            result = y0[:] + (y1[:] - y0[:]) * (date.mjd - y0.date.mjd) / (
                y1.date.mjd - y0.date.mjd
            )

        elif method == self.LAGRANGE:

            stop = prev_idx + 1 + order // 2 + order % 2
            start = prev_idx - order // 2 + 1
            if stop >= len(self):
                start -= stop - len(self)
            elif start < 0:
                stop -= start
                start = 0

            # selection of the subset of data, of length 'order' around the desired value
            subset = self[start:stop]
            date_subset = np.array([x.date.mjd for x in subset])

            result = np.zeros(6)

            if len(subset) < order:
                raise ValueError("len(ephem) < order : impossible to interpolate")

            # Everything is on wikipedia
            #        k
            # L(x) = Σ y_j * l_j(x)
            #        j=0
            #
            # l_j(x) = Π (x - x_m) / (x_j - x_m)
            #     0 <= m <= k
            #        m != j
            for j in range(order):
                # This mask is here to enforce the m != j in the lagrange polynomials
                mask = date_subset != date_subset[j]
                l_j = (date.mjd - date_subset[mask]) / (
                    date_subset[j] - date_subset[mask]
                )
                result = result + l_j.prod() * subset[j]

        else:
            raise ValueError("Unkown interpolation method", method)

        orb = ephem[0]

        return orb.__class__(date, result, orb.form, orb.frame, orb.propagator)

    def propagate(self, date):
        """Alias of :py:meth:`interpolate`
        """
        return self.interpolate(date)

    def iter(
        self, *, dates=None, start=None, stop=None, step=None, strict=True, **kwargs
    ):
        """Ephemeris generator based on the data of this one, but with different dates

        Keyword Arguments:
            dates (list of :py:class:`~beyond.dates.date.Date`): Dates from which iterate over
            start (Date or None): Date of the first point
            stop (Date, timedelta or None): Date of the last point
            step (timedelta or None): Step to use during the computation. Use the same step as
                `self` if `None`
            listeners (list of:py:class:`~beyond.orbits.listeners.Listener`):
            strict (bool): If True, the method will return a ValueError if ``start`` or ``stop`` is
                not in the range of the ephemeris. If False, it will take the closest point in each
                case.
        Yield:
            :py:class:`Orbit`:
        Raise:
            ValueError

        There is two ways to use the iter() method.

        If *dates* is defined, it should be an iterable of dates. This could be
        a generator as per :py:meth:`Date.range <beyond.dates.date.Date.range>`, or a list.

        .. code-block:: python

            # Create two successive ranges of dates, with different steps
            dates = list(Date.range(Date(2019, 3, 23), Date(2019, 3, 24), timedelta(minutes=3)))
            dates.extend(Date.range(Date(2019, 3, 24), Date(2019, 3, 25), timedelta(minutes=10), inclusive=True))
            ephem.iter(dates=dates)

        The alternative, is the use of *start*, *stop* and *step* keyword arguments
        which work exactly as :code:`Date.range(start, stop, step, inclusive=True)`

        If one of *start*, *stop* or *step* arguments is set to ``None`` it will keep
        the same property as the generating ephemeris.

        .. code-block:: python

            # In the examples below, we consider the 'ephem' object to be an ephemeris starting on
            # 2017-01-01 00:00:00 UTC and ending and 2017-01-02 00:00:00 UTC (included) with a fixed
            # step of 3 minutes.

            # These two calls will generate exactly the same points starting at 00:00 and ending at
            # 12:00, as 12:02 does not fall on a date included in the original 'ephem' object.
            ephem.iter(stop=Date(2017, 1, 1, 12))
            ephem.iter(stop=Date(2017, 1, 1, 12, 2))

            # Similarly, these calls will generate the same points starting at 12:00 and ending at
            # 00:00, as 11:58 does not fall on date included in the 'ephem' object.
            ephem.iter(start=Date(2017, 1, 1, 11, 58))
            ephem.iter(start=Date(2017, 1, 1, 12))

            # This call will generate an ephemeris, wich is a subpart of the initial one
            ephem.iter(start=Date(2017, 1, 1, 8), stop=Date(2017, 1, 1, 16))
        """

        # To allow for a loose control of the dates we have to compute
        # the real starting date of the iterator

        listeners = kwargs.get("listeners", [])

        self.clear_listeners(listeners)

        if dates:
            for date in dates:
                orb = self.propagate(date)

                # Listeners
                for listen_orb in self.listen(orb, listeners):
                    yield listen_orb

                yield orb
        else:
            real_start = None

            if start is None:
                start = self.start
            elif start < self.start:
                if strict:
                    raise ValueError("Start date not in range")
                else:
                    real_start = self.start

            if stop is None:
                stop = self.stop
            else:
                if isinstance(stop, timedelta):
                    stop = start + stop
                if stop > self.stop:
                    if strict:
                        raise ValueError("Stop date not in range")
                    else:
                        stop = self.stop

            if real_start is not None:
                start = real_start

            if step is None:

                # The step stays the same as the original ephemeris
                for orb in self:

                    if orb.date < start:
                        continue

                    if orb.date > stop:
                        break

                    # Listeners
                    for listen_orb in self.listen(orb, listeners):
                        yield listen_orb

                    # yield a copy of the recorded orbit to avoid later modification
                    # which could have dire consequences
                    yield orb.copy()
            else:
                # create as ephemeris with a different step than the original
                date = start
                while date <= stop:

                    orb = self.propagate(date)

                    # Listeners
                    for listen_orb in self.listen(orb, listeners):
                        yield listen_orb

                    yield orb
                    date += step

    def ephemeris(self, *args, **kwargs):
        """Same as :py:meth:`self.iter() <iter>`

        Implemented to expose the same methods as :py:class:`Orbit`
        """
        return self.iter(*args, **kwargs)

    def ephem(self, *args, **kwargs):
        """Create an Ephem object which is a subset of this one

        Take the same keyword arguments as :py:meth:`ephemeris`

        Return:
            Ephem:
        """

        return self.__class__(self.ephemeris(*args, **kwargs))

    def copy(self, *, form=None, frame=None):  # pragma: no cover
        """Create a deep copy of the ephemeris, and allow frame and form changing
        """
        new = self.ephem()
        if frame:
            new.frame = frame
        if form:
            new.form = form

        return new

    def as_frame(self, name, **kwargs):  # pragma: no cover
        """Register the Ephem object as a frame

        see :py:func:`beyond.frames.frames.orbit2frame` for details of the arguments
        """
        return orbit2frame(name, self, **kwargs)
<commit_msg>logging ephem : details on range when raising exception<commit_after>"""Definition of ephemeris
"""

import numpy as np
from datetime import timedelta

from .listeners import Speaker
from ..frames.frames import orbit2frame

class Ephem(Speaker):
    """This class represents a range of orbits

    It provides several useful interfaces in order to compute through them

    Example:

    .. code-block:: python

        ephem = orb.ephem(Date.now(), timedelta(hours=6), timedelta(minutes=2))
        ephem.frame = 'ITRF'
        ephem.form = 'spherical'
        latitudes = ephem[:,1]
        longitudes = ephem[:,2]
    """

    LINEAR = "linear"
    LAGRANGE = "lagrange"
    DEFAULT_ORDER = 8

    def __init__(self, orbits, method=None, order=None):
        self._orbits = list(sorted(orbits, key=lambda x: x.date))
        self.method = self.LAGRANGE if method is None else method
        self.order = order if isinstance(order, int) else self.DEFAULT_ORDER

    def __iter__(self):
        self._i = -1
        return self

    def __next__(self):
        self._i += 1
        if self._i >= len(self._orbits):
            raise StopIteration
        return self._orbits[self._i]

    def __getitem__(self, index):
        if isinstance(index, (slice, int)):
            # retrieve a defined Orbit object
            return self._orbits[index]
        else:
            # In order to comply with numpy array handling (e.g. column selection)
            return np.array(self._orbits)[index]

    def __len__(self):
        return len(self._orbits)

    @property
    def start(self):
        """Date of the first element
        """
        return self._orbits[0].date

    @property
    def stop(self):
        """Date of the last element
        """
        return self._orbits[-1].date

    @property
    def dates(self):
        """Generator yielding Dates of each Orbit object of the ephem
        """
        return (o.date for o in self)

    # @property
    # def steps(self):
    #     """Time intervals used in the ephemeris
    #     """
    #     return set([self[i].date - self[i + 1].date for i in range(len(self) - 1)])

    @property
    def frame(self):
        """Get the frame of the first point
        """
        return self._orbits[0].frame

    @frame.setter
    def frame(self, frame):  # pragma: no cover
        """Change the frames of all points
        """
        for orb in self:
            orb.frame = frame

    @property
    def form(self):  # pragma: no cover
        """Get the form of the first point
        """
        return self._orbits[0].form

    @form.setter
    def form(self, form):  # pragma: no cover
        """Change the form of all points
        """
        for orb in self:
            orb.form = form

    def interpolate(self, date, method=None, order=None):
        """Interpolate data at a given date

        Args:
            date (Date):
            method (str): Method of interpolation to use
            order (int): In case of ``LAGRANGE`` method is used
        Return:
            Orbit:
        Raise:
            ValueError: when date is not in the range of the ephemeris
            ValueError: when the order of interpolation is insufficient
        """

        if not self.start <= date <= self.stop:
            raise ValueError("Date '{}' not in range [{}, {}]".format(date, self.start, self.stop))

        prev_idx = 0
        ephem = self

        # Binary search of the orbit step just before the desired date
        while True:
            idx = len(ephem)
            if idx == 1:
                break
            k = idx // 2

            if date > ephem[k].date:
                prev_idx += k
                ephem = ephem[k:]
            else:
                ephem = ephem[:k]

        method = method if method is not None else self.method
        order = order if order is not None else self.order

        if method == self.LINEAR:

            y0 = self[prev_idx]
            y1 = self[prev_idx + 1]

            result = y0[:] + (y1[:] - y0[:]) * (date.mjd - y0.date.mjd) / (
                y1.date.mjd - y0.date.mjd
            )

        elif method == self.LAGRANGE:

            stop = prev_idx + 1 + order // 2 + order % 2
            start = prev_idx - order // 2 + 1
            if stop >= len(self):
                start -= stop - len(self)
            elif start < 0:
                stop -= start
                start = 0

            # selection of the subset of data, of length 'order' around the desired value
            subset = self[start:stop]
            date_subset = np.array([x.date.mjd for x in subset])

            result = np.zeros(6)

            if len(subset) < order:
                raise ValueError("len(ephem) < order : impossible to interpolate")

            # Everything is on wikipedia
            #        k
            # L(x) = Σ y_j * l_j(x)
            #        j=0
            #
            # l_j(x) = Π (x - x_m) / (x_j - x_m)
            #     0 <= m <= k
            #        m != j
            for j in range(order):
                # This mask is here to enforce the m != j in the lagrange polynomials
                mask = date_subset != date_subset[j]
                l_j = (date.mjd - date_subset[mask]) / (
                    date_subset[j] - date_subset[mask]
                )
                result = result + l_j.prod() * subset[j]

        else:
            raise ValueError("Unkown interpolation method", method)

        orb = ephem[0]

        return orb.__class__(date, result, orb.form, orb.frame, orb.propagator)

    def propagate(self, date):
        """Alias of :py:meth:`interpolate`
        """
        return self.interpolate(date)

    def iter(
        self, *, dates=None, start=None, stop=None, step=None, strict=True, **kwargs
    ):
        """Ephemeris generator based on the data of this one, but with different dates

        Keyword Arguments:
            dates (list of :py:class:`~beyond.dates.date.Date`): Dates from which iterate over
            start (Date or None): Date of the first point
            stop (Date, timedelta or None): Date of the last point
            step (timedelta or None): Step to use during the computation. Use the same step as
                `self` if `None`
            listeners (list of:py:class:`~beyond.orbits.listeners.Listener`):
            strict (bool): If True, the method will return a ValueError if ``start`` or ``stop`` is
                not in the range of the ephemeris. If False, it will take the closest point in each
                case.
        Yield:
            :py:class:`Orbit`:
        Raise:
            ValueError

        There is two ways to use the iter() method.

        If *dates* is defined, it should be an iterable of dates. This could be
        a generator as per :py:meth:`Date.range <beyond.dates.date.Date.range>`, or a list.

        .. code-block:: python

            # Create two successive ranges of dates, with different steps
            dates = list(Date.range(Date(2019, 3, 23), Date(2019, 3, 24), timedelta(minutes=3)))
            dates.extend(Date.range(Date(2019, 3, 24), Date(2019, 3, 25), timedelta(minutes=10), inclusive=True))
            ephem.iter(dates=dates)

        The alternative, is the use of *start*, *stop* and *step* keyword arguments
        which work exactly as :code:`Date.range(start, stop, step, inclusive=True)`

        If one of *start*, *stop* or *step* arguments is set to ``None`` it will keep
        the same property as the generating ephemeris.

        .. code-block:: python

            # In the examples below, we consider the 'ephem' object to be an ephemeris starting on
            # 2017-01-01 00:00:00 UTC and ending and 2017-01-02 00:00:00 UTC (included) with a fixed
            # step of 3 minutes.

            # These two calls will generate exactly the same points starting at 00:00 and ending at
            # 12:00, as 12:02 does not fall on a date included in the original 'ephem' object.
            ephem.iter(stop=Date(2017, 1, 1, 12))
            ephem.iter(stop=Date(2017, 1, 1, 12, 2))

            # Similarly, these calls will generate the same points starting at 12:00 and ending at
            # 00:00, as 11:58 does not fall on date included in the 'ephem' object.
            ephem.iter(start=Date(2017, 1, 1, 11, 58))
            ephem.iter(start=Date(2017, 1, 1, 12))

            # This call will generate an ephemeris, wich is a subpart of the initial one
            ephem.iter(start=Date(2017, 1, 1, 8), stop=Date(2017, 1, 1, 16))
        """

        # To allow for a loose control of the dates we have to compute
        # the real starting date of the iterator

        listeners = kwargs.get("listeners", [])

        self.clear_listeners(listeners)

        if dates:
            for date in dates:
                orb = self.propagate(date)

                # Listeners
                for listen_orb in self.listen(orb, listeners):
                    yield listen_orb

                yield orb
        else:
            real_start = None

            if start is None:
                start = self.start
            elif start < self.start:
                if strict:
                    raise ValueError("Start date not in range [{}, {}]".format(self.start, self.stop))
                else:
                    real_start = self.start

            if stop is None:
                stop = self.stop
            else:
                if isinstance(stop, timedelta):
                    stop = start + stop
                if stop > self.stop:
                    if strict:
                        raise ValueError("Stop date not in range [{}, {}]".format(self.start, self.stop))
                    else:
                        stop = self.stop

            if real_start is not None:
                start = real_start

            if step is None:

                # The step stays the same as the original ephemeris
                for orb in self:

                    if orb.date < start:
                        continue

                    if orb.date > stop:
                        break

                    # Listeners
                    for listen_orb in self.listen(orb, listeners):
                        yield listen_orb

                    # yield a copy of the recorded orbit to avoid later modification
                    # which could have dire consequences
                    yield orb.copy()
            else:
                # create as ephemeris with a different step than the original
                date = start
                while date <= stop:

                    orb = self.propagate(date)

                    # Listeners
                    for listen_orb in self.listen(orb, listeners):
                        yield listen_orb

                    yield orb
                    date += step

    def ephemeris(self, *args, **kwargs):
        """Same as :py:meth:`self.iter() <iter>`

        Implemented to expose the same methods as :py:class:`Orbit`
        """
        return self.iter(*args, **kwargs)

    def ephem(self, *args, **kwargs):
        """Create an Ephem object which is a subset of this one

        Take the same keyword arguments as :py:meth:`ephemeris`

        Return:
            Ephem:
        """

        return self.__class__(self.ephemeris(*args, **kwargs))

    def copy(self, *, form=None, frame=None):  # pragma: no cover
        """Create a deep copy of the ephemeris, and allow frame and form changing
        """
        new = self.ephem()
        if frame:
            new.frame = frame
        if form:
            new.form = form

        return new

    def as_frame(self, name, **kwargs):  # pragma: no cover
        """Register the Ephem object as a frame

        see :py:func:`beyond.frames.frames.orbit2frame` for details of the arguments
        """
        return orbit2frame(name, self, **kwargs)
<|endoftext|>


<commit_before># Copyright 2019 PerfKitBenchmarker Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Module containing class for AWS's Athena EDW service."""

from perfkitbenchmarker import edw_service
from perfkitbenchmarker import flags
from perfkitbenchmarker import providers

FLAGS = flags.FLAGS

class Athena(edw_service.EdwService):
  """Object representing a Athena data warehouse."""

  CLOUD = providers.AWS
  SERVICE_TYPE = 'athena'

  def _Create(self):
    """Create a Athena data warehouse."""
    raise NotImplementedError

  def _Exists(self):
    """Method to validate the existence of a Athena data warehouse.

    Returns:
      Boolean value indicating the existence of a Athena data warehouse.
    """
    raise NotImplementedError

  def _Delete(self):
    """Delete a Athena data warehouse."""
    raise NotImplementedError

  def GetMetadata(self):
    """Return a dictionary of the metadata for the Athena data warehouse."""
    basic_data = super(Athena, self).GetMetadata()
    return basic_data

  def RunCommandHelper(self):
    """Athena data warehouse specific run script command components."""
    raise NotImplementedError

  def InstallAndAuthenticateRunner(self, vm):
    """Method to perform installation and authentication of Athena query runner.

    Athena APIs are included in the default AWS cli
    https://docs.aws.amazon.com/cli/latest/reference/athena/index.html.

    Args:
      vm: Client vm on which the query will be run.
    """
    for pkg in ('aws_credentials', 'awscli'):
      vm.Install(pkg)

  def PrepareClientVm(self, vm):
    """Prepare phase to install the runtime environment on the client vm.

    Args:
      vm: Client vm on which the script will be run.
    """
    super(Athena, self).PrepareClientVm(vm)
    self.InstallAndAuthenticateRunner(vm)

  @classmethod
  def RunScriptOnClientVm(cls, vm, database, script):
    """A function to execute a script against AWS Athena Database."""
    stdout, _ = vm.RemoteCommand('python runner.py --database=%s --script=%s' %
                                 (database, script))
    return stdout
<commit_msg>AWS Athena provisioning functionality.<commit_after># Copyright 2019 PerfKitBenchmarker Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Module containing class for AWS's Athena EDW service."""

import datetime
import logging

from perfkitbenchmarker import data

from perfkitbenchmarker import edw_service
from perfkitbenchmarker import flags
from perfkitbenchmarker import providers
from perfkitbenchmarker import vm_util
from perfkitbenchmarker.providers.aws import util

AWS_ATHENA_CMD_PREFIX = ['aws', 'athena']
AWS_ATHENA_CMD_POSTFIX = ['--output', 'json']

FLAGS = flags.FLAGS

class AthenaQueryError(RuntimeError):
  pass

def ReadScript(script_uri):
  """Method to read a sql script based on its local path.

  Arguments:
    script_uri: Local URI of file containing SQL query.

  Returns:
    Query String contents of the URI location.

  Raises:
    IOError: If the script cannot be read.
  """
  with open(script_uri) as fp:
    return fp.read()

def PrepareQueryString(query_string_template, substitutions):
  """Method to read a template Athena script and substitute placeholders.

  Args:
    query_string_template: Template version of the Athena query.
    substitutions: A dictionary of string placeholder keys and corresponding
      string values.

  Returns:
     Materialized Athena query as a string.
  """
  for key, value in substitutions.items():
    query_string = query_string_template.replace(key, value)
  return query_string

def RunScriptCommand(script_command):
  """Method to execute an AWS Athena cli command.

  Args:
    script_command: Fully compiled AWS Athena cli command.

  Returns:
    String stdout result of executing the query.
    Script Command execution duration in seconds (rounded).

  Raises:
    AthenaQueryError: If the return code does not indicate success.
  """
  start_time = datetime.datetime.now()
  stdout, _, retcode = vm_util.IssueCommand(
      script_command, raise_on_failure=False)
  if retcode:
    raise AthenaQueryError
  end_time = datetime.datetime.now()
  return stdout, int((end_time - start_time).total_seconds())

class Athena(edw_service.EdwService):
  """Object representing a Athena data warehouse."""

  CLOUD = providers.AWS
  SERVICE_TYPE = 'athena'

  def __init__(self, edw_service_spec):
    super(Athena, self).__init__(edw_service_spec)
    self.output_location = FLAGS.athena_output_location
    self.region = util.GetRegionFromZone(FLAGS.zones[0])
    if FLAGS.provision_athena:
      self.db = '_'.join(
          [FLAGS.edw_tpc_dsb_type,
           str(FLAGS.edw_tpc_dataset_size_in_GB)])
      self.data_bucket = self.db.replace('_', '')
      self.athena_db_create_time = 0
      self.athena_table_create_time = 0

  def BuildAthenaCommand(self, query_string, database=None):
    """Method to compile a AWS Athena cli command.

    Arguments:
      query_string: A string with the query that needs to be executed on Athena.
      database: The Athena database against which the query should be executed.

    Returns:
      Fully compiled AWS Athena cli command.
    """
    cmd = []
    cmd.extend(AWS_ATHENA_CMD_PREFIX)
    cmd.extend([
        '--region', self.region,
        'start-query-execution',
        '--query-string', query_string
    ])
    if database:
      cmd.extend(['--query-execution-context', ('Database=%s' % database)])
    cmd.extend([
        '--result-configuration', ('OutputLocation=%s' % self.output_location)
    ])
    cmd.extend(AWS_ATHENA_CMD_POSTFIX)
    return cmd

  def _Create(self):
    """Create a Athena data warehouse."""

    def _CreateDatabase():
      create_database_query_string = PrepareQueryString(
          'create database database_name', {'database_name': self.db})
      script_command = self.BuildAthenaCommand(create_database_query_string)
      return RunScriptCommand(script_command)

    def _CreateTable(table_create_sql_template):
      template_script_path = data.ResourcePath(table_create_sql_template)
      template_script_contents = ReadScript(template_script_path)
      script_contents = PrepareQueryString(template_script_contents,
                                           {'{bucket}': self.data_bucket})
      script_command = self.BuildAthenaCommand(
          script_contents, database=self.db)
      return RunScriptCommand(script_command)

    def _CreateAllTables():
      cumulative_table_create_time = 0
      # TODO(saksena): Derive the paths and full table set from the TPC suite.
      for script in [
          'edw/athena/tpc_h/ddl/s3_customer.sql',
          'edw/athena/tpc_h/ddl/s3_nation.sql'
      ]:
        _, table_create_time = _CreateTable(script)
        cumulative_table_create_time += table_create_time
      return cumulative_table_create_time

    _, self.athena_db_create_time = _CreateDatabase()
    self.athena_table_create_time = _CreateAllTables()

  def _Exists(self):
    """Method to validate the existence of a Athena data warehouse.

    Returns:
      Boolean value indicating the existence of a Athena data warehouse.
    """
    raise NotImplementedError

  def _Delete(self):
    """Delete a Athena data warehouse."""
    if not FLAGS.teardown_athena:
      logging.info('The current resource is requested to be long living.')
      return
    raise NotImplementedError

  def GetMetadata(self):
    """Return a dictionary of the metadata for the Athena data warehouse."""
    basic_data = super(Athena, self).GetMetadata()
    return basic_data

  def RunCommandHelper(self):
    """Athena data warehouse specific run script command components."""
    raise NotImplementedError

  def InstallAndAuthenticateRunner(self, vm):
    """Method to perform installation and authentication of Athena query runner.

    Athena APIs are included in the default AWS cli
    https://docs.aws.amazon.com/cli/latest/reference/athena/index.html.

    Args:
      vm: Client vm on which the query will be run.
    """
    for pkg in ('aws_credentials', 'awscli'):
      vm.Install(pkg)

  def PrepareClientVm(self, vm):
    """Prepare phase to install the runtime environment on the client vm.

    Args:
      vm: Client vm on which the script will be run.
    """
    super(Athena, self).PrepareClientVm(vm)
    self.InstallAndAuthenticateRunner(vm)

  @classmethod
  def RunScriptOnClientVm(cls, vm, database, script):
    """A function to execute a script against AWS Athena Database."""
    stdout, _ = vm.RemoteCommand('python runner.py --database=%s --script=%s' %
                                 (database, script))
    return stdout
<|endoftext|>


<commit_before># plot the Hugoniot loci for a compressible Riemann problem

import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as optimize

class State(object):
    def __init__(self, p=1.0, u=0.0, rho=1.0):
         self.p = p
         self.u = u
         self.rho = rho

    def __str__(self):
        return "rho: {}; u: {}; p: {}".format(self.rho, self.u, self.p)

class RiemannProblem(object):
    def __init__(self, left_state, right_state, gamma=1.4):
        self.left = left_state
        self.right = right_state
        self.gamma = gamma

    def u_hugoniot(self, p, side):

        if side == "left":
            state = self.left
            s = 1.0
        elif side == "right":
            state = self.right
            s = -1.0

        c = np.sqrt(self.gamma*state.p/state.rho)

        if p < state.p:
            # rarefaction
            u = state.u + s*(2.0*c/(self.gamma-1.0))* \
                (1.0 - (p/state.p)**((self.gamma-1.0)/(2.0*self.gamma)))
        else:
            # shock
            beta = (self.gamma+1.0)/(self.gamma-1.0)
            u = state.u + s*(2.0*c/np.sqrt(2.0*self.gamma*(self.gamma-1.0)))* \
                (1.0 - p/state.p)/np.sqrt(1.0 + beta*p/state.p)

        return u

    def find_star_state(self, p_min=0.001, p_max=1000.0):
        # we need to root-find on
        pstar = optimize.brentq(lambda p: self.u_hugoniot(p, "left") - self.u_hugoniot(p, "right"),
                               p_min, p_max)
        ustar = self.u_hugoniot(pstar, "left")

        return pstar, ustar

    def plot_hugoniot(self, p_min = 0.0, p_max=1.5, N=200):

        p = np.linspace(p_min, p_max, num=N)
        u_left = np.zeros_like(p)
        u_right = np.zeros_like(p)

        for n in range(N):
            u_left[n] = self.u_hugoniot(p[n], "left")
        ish = np.where(p > self.left.p)
        ir = np.where(p < self.left.p)

        plt.plot(p[ish], u_left[ish], c="b", ls=":", lw=2)
        plt.plot(p[ir], u_left[ir], c="b", ls="-", lw=2)
        plt.scatter([self.left.p], [self.left.u], marker="x", c="b", s=40)

        for n in range(N):
            u_right[n] = self.u_hugoniot(p[n], "right")
        ish = np.where(p > self.right.p)
        ir = np.where(p < self.right.p)

        plt.plot(p[ish], u_right[ish], c="r", ls=":", lw=2)
        plt.plot(p[ir], u_right[ir], c="r", ls="-", lw=2)
        plt.scatter([self.right.p], [self.right.u], marker="x", c="r", s=40)

        du = 0.025*(max(np.max(u_left), np.max(u_right)) -
                    min(np.min(u_left), np.min(u_right)))

        plt.text(self.left.p, self.left.u+du, "left",
                 horizontalalignment="center", color="b")

        plt.text(self.right.p, self.right.u+du, "right",
                 horizontalalignment="center", color="r")

        plt.xlim(p_min, p_max)

        plt.xlabel(r"$p$", fontsize="large")
        plt.ylabel(r"$u$", fontsize="large")

        legs = []
        legnames = []

        legs.append(plt.Line2D((0,1),(0,0), color="k", ls=":", marker=None))
        legnames.append("shock")

        legs.append(plt.Line2D((0,1),(0,0), color="k", ls="-", marker=None))
        legnames.append("rarefaction")

        plt.legend(legs, legnames, frameon=False, loc="best")

        plt.tight_layout()

        plt.savefig("riemann-phase.pdf")

if __name__ == "__main__":

    # setup the problem -- Sod
    left = State(p = 1.0, u = 0.0, rho = 1.0)
    right = State(p = 0.1, u = 0.0, rho = 0.125)

    rp = RiemannProblem(left, right)

    rp.plot_hugoniot()

    pstar, ustar = rp.find_star_state()

    print pstar, ustar
<commit_msg>fix how we mark shock vs rarefaction<commit_after># plot the Hugoniot loci for a compressible Riemann problem

import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as optimize

class State(object):
    def __init__(self, p=1.0, u=0.0, rho=1.0):
         self.p = p
         self.u = u
         self.rho = rho

    def __str__(self):
        return "rho: {}; u: {}; p: {}".format(self.rho, self.u, self.p)

class RiemannProblem(object):
    def __init__(self, left_state, right_state, gamma=1.4):
        self.left = left_state
        self.right = right_state
        self.gamma = gamma

    def u_hugoniot(self, p, side):

        if side == "left":
            state = self.left
            s = 1.0
        elif side == "right":
            state = self.right
            s = -1.0

        c = np.sqrt(self.gamma*state.p/state.rho)

        if p < state.p:
            # rarefaction
            u = state.u + s*(2.0*c/(self.gamma-1.0))* \
                (1.0 - (p/state.p)**((self.gamma-1.0)/(2.0*self.gamma)))
        else:
            # shock
            beta = (self.gamma+1.0)/(self.gamma-1.0)
            u = state.u + s*(2.0*c/np.sqrt(2.0*self.gamma*(self.gamma-1.0)))* \
                (1.0 - p/state.p)/np.sqrt(1.0 + beta*p/state.p)

        return u

    def find_star_state(self, p_min=0.001, p_max=1000.0):
        # we need to root-find on
        pstar = optimize.brentq(lambda p: self.u_hugoniot(p, "left") - self.u_hugoniot(p, "right"),
                               p_min, p_max)
        ustar = self.u_hugoniot(pstar, "left")

        return pstar, ustar

    def plot_hugoniot(self, pstar, ustar, p_min = 0.0, p_max=1.5, N=200):

        p = np.linspace(p_min, p_max, num=N)
        u_left = np.zeros_like(p)
        u_right = np.zeros_like(p)

        for n in range(N):
            u_left[n] = self.u_hugoniot(p[n], "left")
        ish = np.where(p > pstar)
        ir = np.where(p < pstar)

        plt.plot(p[ish], u_left[ish], c="b", ls=":", lw=2)
        plt.plot(p[ir], u_left[ir], c="b", ls="-", lw=2)
        plt.scatter([self.left.p], [self.left.u], marker="x", c="b", s=40)

        for n in range(N):
            u_right[n] = self.u_hugoniot(p[n], "right")
        ish = np.where(p > pstar)
        ir = np.where(p < pstar)

        plt.plot(p[ish], u_right[ish], c="r", ls=":", lw=2)
        plt.plot(p[ir], u_right[ir], c="r", ls="-", lw=2)
        plt.scatter([self.right.p], [self.right.u], marker="x", c="r", s=40)

        du = 0.025*(max(np.max(u_left), np.max(u_right)) -
                    min(np.min(u_left), np.min(u_right)))

        plt.text(self.left.p, self.left.u+du, "left",
                 horizontalalignment="center", color="b")

        plt.text(self.right.p, self.right.u+du, "right",
                 horizontalalignment="center", color="r")

        plt.xlim(p_min, p_max)

        plt.xlabel(r"$p$", fontsize="large")
        plt.ylabel(r"$u$", fontsize="large")

        legs = []
        legnames = []

        legs.append(plt.Line2D((0,1),(0,0), color="k", ls=":", marker=None))
        legnames.append("shock")

        legs.append(plt.Line2D((0,1),(0,0), color="k", ls="-", marker=None))
        legnames.append("rarefaction")

        plt.legend(legs, legnames, frameon=False, loc="best")

        plt.tight_layout()

        plt.savefig("riemann-phase.pdf")

if __name__ == "__main__":

    # setup the problem -- Sod
    left = State(p = 1.0, u = 0.0, rho = 1.0)
    right = State(p = 0.1, u = 0.0, rho = 0.125)

    rp = RiemannProblem(left, right)

    pstar, ustar = rp.find_star_state()

    rp.plot_hugoniot(pstar, ustar)

    print pstar, ustar
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
Pull together the NCES data and the Segragation index calculator to produce
some reports on segregation in the USA.
"""
import sys
import argparse

from segcalc import SegCalc
from nces_parser import NCESParser

from xlwt import Workbook

# ==============================================================================
# Constants
# ==============================================================================
fips_to_st = {
    1: ('AL', 'Alabama'),
    2: ('AK', 'Alaska'),
    4: ('AZ', 'Arizona'),
    5: ('AR', 'Arkansas'),
    6: ('CA', 'California'),
    8: ('CO', 'Colorado'),
    9: ('CT', 'Connecticut'),
    10: ('DE', 'Delaware'),
    11: ('DC', 'District of Columbia'),
    12: ('FL', 'Florida'),
    13: ('GA', 'Georgia'),
    15: ('HI', 'Hawaii'),
    16: ('ID', 'Idaho'),
    17: ('IL', 'Illinois'),
    18: ('IN', 'Indiana'),
    19: ('IA', 'Iowa'),
    20: ('KS', 'Kansas'),
    21: ('KY', 'Kentucky'),
    22: ('LA', 'Louisiana'),
    23: ('ME', 'Maine'),
    24: ('MD', 'Maryland'),
    25: ('MA', 'Massachusetts'),
    26: ('MI', 'Michigan'),
    27: ('MN', 'Minnesota'),
    28: ('MS', 'Mississippi'),
    29: ('MO', 'Missouri'),
    30: ('MT', 'Montana'),
    31: ('NE', 'Nebraska'),
    32: ('NV', 'Nevada'),
    33: ('NH', 'New Hampshire'),
    34: ('NJ', 'New Jersey'),
    35: ('NM', 'New Mexico'),
    36: ('NY', 'New York'),
    37: ('NC', 'North Carolina'),
    38: ('ND', 'North Dakota'),
    39: ('OH', 'Ohio'),
    40: ('OK', 'Oklahoma'),
    41: ('OR', 'Oregon'),
    42: ('PA', 'Pennsylvania'),
    44: ('RI', 'Rhode Island'),
    45: ('SC', 'South Carolina'),
    46: ('SD', 'South Dakota'),
    47: ('TN', 'Tennessee'),
    48: ('TX', 'Texas'),
    49: ('UT', 'Utah'),
    50: ('VT', 'Vermont'),
    51: ('VA', 'Virginia'),
    53: ('WA', 'Washington'),
    54: ('WV', 'West Virginia'),
    55: ('WI', 'Wisconsin'),
    56: ('WY', 'Wyoming')
}

# ==============================================================================
# Functions
# ==============================================================================
def calc_idxes(year_range, idx):
    """
    Call down to get all the various measures calculated
    """
    exp_idx = []
    iso_idx = []
    dis_idx = []
    for year in year_range:
        print "Working on NCES Data from:  %d" % year
        # Get our dataset ready
        nces = NCESParser(year=year)
        schools = nces.parse(make_dict=True)
        sg = SegCalc(schools, idx)
        exp_idx.append(sg.calc_exp_idx())
        iso_idx.append(sg.calc_iso_idx())
        dis_idx.append(sg.calc_dis_idx())
    return (exp_idx, iso_idx, dis_idx)

# -------------------------------------
def save_report(year_range, idxes, filename):
    # Break out the calculated data
    exp_idx, iso_idx, dis_idx = idxes

    # Now lets format it into a report.
    # Report will be a 2D matrix
    # X-Axis = school year
    # Y-Axis = FIPS Code
    wb = Workbook()
    ews = wb.add_sheet('Exposure Index')
    iws = wb.add_sheet('Isolation Index')
    dws = wb.add_sheet('Dissimilarity Index')

    worksheets = [ews, iws, dws]

    # Create the headers/labels row/col
    for ws in worksheets:
        ws.write(0, 0, "State/Year")
        for j, st in enumerate(fips_to_st):
            ws.write(j+1, 0, fips_to_st[st][0])

    # Print out the data
    for i, year in enumerate(year_range):
        print "Write Report for:  %d" % year
        for ws in worksheets:
            ws.write(0, i+1, year)
        for j, st in enumerate(fips_to_st.keys()):
            if exp_idx[i][st] < 0.001:
                ews.write(j+1, i+1, "")
            else:
                ews.write(j+1, i+1, exp_idx[i][st])

            if iso_idx[i][st] < 0.001:
                iws.write(j+1, i+1, "")
            else:
                iws.write(j+1, i+1, iso_idx[i][st])

            if dis_idx[i][st] < 0.001:
                dws.write(j+1, i+1, "")
            else:
                dws.write(j+1, i+1, dis_idx[i][st])
    wb.save(filename)

# -------------------------------------
# Parse the command line options
# -------------------------------------
def main(argv):
    parser = argparse.ArgumentParser(description='Segregation Report Generator')
    parser.add_argument('--year', action='store', dest='year', required=False, type=int,
            help='NCES Data Year')
    parser.add_argument('--outfile', action='store', dest='outfile', required=False,
            help='Report File')
    args = parser.parse_args()

    # Lets calculate all the data first
    year_range = range(1987, 2011)

    """
    idx = {'Y_GROUP': 'BLACK', 'Z_GROUP': 'WHITE', 'TOTAL': 'MEMBER', 'CATEGORY': 'FIPS', 'SUB_CAT': 'LEAID'}
    idxes = calc_idxes(year_range, idx)
    save_report(year_range, idxes, "black_" + args.outfile)

    idx = {'Y_GROUP': 'HISP', 'Z_GROUP': 'WHITE', 'TOTAL': 'MEMBER', 'CATEGORY': 'FIPS', 'SUB_CAT': 'LEAID'}
    idxes = calc_idxes(year_range, idx)
    save_report(year_range, idxes, "hisp_" + args.outfile)

    idx = {'Y_GROUP': 'ASIAN', 'Z_GROUP': 'WHITE', 'TOTAL': 'MEMBER', 'CATEGORY': 'FIPS', 'SUB_CAT': 'LEAID'}
    idxes = calc_idxes(year_range, idx)
    save_report(year_range, idxes, "asian_" + args.outfile)
    """

    idx = {'Y_GROUP': 'AM', 'Z_GROUP': 'WHITE', 'TOTAL': 'MEMBER', 'CATEGORY': 'FIPS', 'SUB_CAT': 'LEAID'}
    idxes = calc_idxes(year_range, idx)
    save_report(year_range, idxes, "native_amer_" + args.outfile)

# -------------------------------------
# Drop the script name from the args
# and call our command line parser
# -------------------------------------
if __name__ == "__main__":
    main(sys.argv[1:])

<commit_msg>Modified to allow a broader set of analysis options:<commit_after>#!/usr/bin/env python
"""
Pull together the NCES data and the Segragation index calculator to produce
some reports on segregation in the USA.
"""
import sys
import argparse

from segcalc import SegCalc
from nces_parser import NCESParser

from xlwt import Workbook

# ==============================================================================
# Constants
# ==============================================================================
fips_to_st = {
    1: ('AL', 'Alabama'),
    2: ('AK', 'Alaska'),
    4: ('AZ', 'Arizona'),
    5: ('AR', 'Arkansas'),
    6: ('CA', 'California'),
    8: ('CO', 'Colorado'),
    9: ('CT', 'Connecticut'),
    10: ('DE', 'Delaware'),
    11: ('DC', 'District of Columbia'),
    12: ('FL', 'Florida'),
    13: ('GA', 'Georgia'),
    15: ('HI', 'Hawaii'),
    16: ('ID', 'Idaho'),
    17: ('IL', 'Illinois'),
    18: ('IN', 'Indiana'),
    19: ('IA', 'Iowa'),
    20: ('KS', 'Kansas'),
    21: ('KY', 'Kentucky'),
    22: ('LA', 'Louisiana'),
    23: ('ME', 'Maine'),
    24: ('MD', 'Maryland'),
    25: ('MA', 'Massachusetts'),
    26: ('MI', 'Michigan'),
    27: ('MN', 'Minnesota'),
    28: ('MS', 'Mississippi'),
    29: ('MO', 'Missouri'),
    30: ('MT', 'Montana'),
    31: ('NE', 'Nebraska'),
    32: ('NV', 'Nevada'),
    33: ('NH', 'New Hampshire'),
    34: ('NJ', 'New Jersey'),
    35: ('NM', 'New Mexico'),
    36: ('NY', 'New York'),
    37: ('NC', 'North Carolina'),
    38: ('ND', 'North Dakota'),
    39: ('OH', 'Ohio'),
    40: ('OK', 'Oklahoma'),
    41: ('OR', 'Oregon'),
    42: ('PA', 'Pennsylvania'),
    44: ('RI', 'Rhode Island'),
    45: ('SC', 'South Carolina'),
    46: ('SD', 'South Dakota'),
    47: ('TN', 'Tennessee'),
    48: ('TX', 'Texas'),
    49: ('UT', 'Utah'),
    50: ('VT', 'Vermont'),
    51: ('VA', 'Virginia'),
    53: ('WA', 'Washington'),
    54: ('WV', 'West Virginia'),
    55: ('WI', 'Wisconsin'),
    56: ('WY', 'Wyoming')
}

# Maximum number of rows to report in an output file.
MAX_REPORT = 200

# ==============================================================================
# Functions
# ==============================================================================
def calc_idxes(segcalc):
    """
    Call down to get all the various measures calculated
    """
    exp_idx = segcalc.calc_exp_idx()
    iso_idx = segcalc.calc_iso_idx()
    dis_idx = segcalc.calc_dis_idx()
    tot_idx = segcalc.calc_totals()
    return (exp_idx, iso_idx, dis_idx, tot_idx)

# --------------------------------------
def calc_idxes_range(year_range, idx):
    """
    Get all the results over the years
    """
    exp_idx = []
    iso_idx = []
    dis_idx = []
    for year in year_range:
        print "Working on NCES Data from:  %d" % year
        # Get our dataset ready
        nces = NCESParser(year=year)
        schools = nces.parse(make_dict=True)
        sg = SegCalc(schools, idx)
        exp_idx.append(sg.calc_exp_idx())
        iso_idx.append(sg.calc_iso_idx())
        dis_idx.append(sg.calc_dis_idx())
    return (exp_idx, iso_idx, dis_idx)

# -------------------------------------
def save_report(year_range, idxes, category_list, category_txt, filename):
    """
    Write out a bunch of report data to a spreadsheet report.
    Report will be a 2D matrix:
        - X-Axis = school year
        - Y-Axis = 'Category' (FIPS Code, District, etc...)

    Notes:
        - idxes contains the data
        - worksheets is a list of XLS worksheets, one per report in idxes
    """
    wb = Workbook()
    ews = wb.add_sheet('Exposure Index')
    iws = wb.add_sheet('Isolation Index')
    dws = wb.add_sheet('Dissimilarity Index')
    size = wb.add_sheet('Student Count')

    worksheets = [ews, iws, dws, size]

    # Create the headers/labels row/col
    for ws in worksheets:
        ws.write(0, 0, "LEA/Year")
        for j, st in enumerate(category_list):
            if j < MAX_REPORT:
                if len(category_txt[st]) == 2: # Don't change caps for State abbr.
                    ws.write(j+1, 0, category_txt[st])
                else:
                    ws.write(j+1, 0, category_txt[st].title())

    # Print out the data
    for i, year in enumerate(year_range):
        print "Write Report for:  %d" % year
        for ws in worksheets:
            ws.write(0, i+1, year)
        for j, st in enumerate(category_list):
            if j < MAX_REPORT:
                for k, idx in enumerate(idxes):
                    try:
                        if idx[i][st] < 0.001:
                            worksheets[k].write(j+1, i+1, "")
                        else:
                            worksheets[k].write(j+1, i+1, idx[i][st])
                    except KeyError:
                        worksheets[k].write(j+1, i+1, "")
    wb.save(filename)

# -------------------------------------
# Parse the command line options
# -------------------------------------
def main(argv):
    parser = argparse.ArgumentParser(description='Segregation Report Generator')
    parser.add_argument('--year', action='store', dest='year', required=False, type=int,
            help='NCES Data Year')
    parser.add_argument('--outfile', action='store', dest='outfile', required=True,
            help='Report File')
    parser.add_argument('--category', action='store', dest='category', required=False,
            help='Which Category do we sort the results by?')
    parser.add_argument('--match_idx', action='store', dest='match_idx', required=False,
            help='Only use data points that match some criterion')
    parser.add_argument('--match_val', action='store', dest='match_val', required=False,
            help='Only use data points that match some criterion')
    parser.add_argument('-debug', action='store_true', dest='debug', required=False,
            help='Debug Mode')
    args = parser.parse_args()

    if args.category:
        category = args.category
    else:
        category = 'FIPS'

    # Lets calculate all the data first
    if args.debug:
        year_range = [1987]
        groups = ['BLACK']
    else:
        year_range = range(1987, 2011)
        groups = ['BLACK', 'HISP', 'ASIAN', 'AM']

    # Default search query
    idx = {
        'Y_GROUP': "BLACK",
        'Z_GROUP': 'WHITE',
        'TOTAL': 'MEMBER',
        'CATEGORY': category,
        'SUB_CAT': 'LEAID',
    }
    if args.match_idx:
        idx['MATCH_IDX'] = args.match_idx
        idx['MATCH_VAL'] = args.match_val

    idxes = [[], [], [], []]
    for group in groups:
        for year in year_range:
            print "Loading NCES Data from:  %d" % year
            nces = NCESParser(year=year)
            schools = nces.parse(make_dict=True)
            # Get our data query ready
            idx['Y_GROUP'] = group
            segcalc = SegCalc(schools, idx)
            if category == 'LEAID':
                category_lut = segcalc.get_idxed_val('LEAID', 'LEANM')
            elif category == 'FIPS':
                category_lut = dict(zip(fips_to_st.keys(), [fips_to_st[key][0] for key in fips_to_st.keys()]))

            exp,iso,dis,tot = calc_idxes(segcalc)
            category_by_size = sorted(tot, key=tot.get, reverse=True)
            # Filter out keys absent from our report tables.
            category_by_size = [i for i in category_by_size if i in category_lut.keys()]

            idxes[0].append(exp)
            idxes[1].append(iso)
            idxes[2].append(dis)
            idxes[3].append(tot)

        save_report(year_range, idxes, category_by_size, category_lut, group.lower() + '_' + args.outfile)

# -------------------------------------
# Drop the script name from the args
# and call our command line parser
# -------------------------------------
if __name__ == "__main__":
    main(sys.argv[1:])

<|endoftext|>


<commit_before># Copyright (c) 2006-2007 rPath, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.opensource.org/licenses/cpl.php.
#
# This program 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 Common Public License for
# full details.

from conary import files, trove, versions
from conary import errors as conaryerrors
from conary.build import build, source
from conary.build import errors as builderrors
from conary.build.packagerecipe import _AbstractPackageRecipe
from conary.local import update
from conary.lib import log, util
from conary.repository import changeset, filecontents

class DerivedPackageRecipe(_AbstractPackageRecipe):

    internalAbstractBaseClass = 1
    _isDerived = True
    parentVersion = None

    def _expandChangeset(self):
        destdir = self.macros.destdir

        delayedRestores = {}
        ptrMap = {}
        byDefault = {}

        fileList = []
        linkGroups = {}
        linkGroupFirstPath = {}
        # sort the files by pathId,fileId
        for trvCs in self.cs.iterNewTroveList():
            trv = trove.Trove(trvCs)

            # these should all be the same anyway
            flavor = trv.getFlavor().copy()
            name = trv.getName()
            self._componentReqs[name] = trv.getRequires().copy()
            self._componentProvs[name] = trv.getProvides().copy()

            for pathId, path, fileId, version in trv.iterFileList():
                if path != self.macros.buildlogpath:
                    fileList.append((pathId, fileId, path, name))

            if trv.isCollection():
                # gather up existing byDefault status
                # from (component, byDefault) tuples
                byDefault.update(dict(
                    [(x[0][0], x[1]) for x in trv.iterTroveListInfo()]))

        fileList.sort()

        for pathId, fileId, path, troveName in fileList:
            fileCs = self.cs.getFileChange(None, fileId)
            fileObj = files.ThawFile(fileCs, pathId)
            self._derivedFiles[path] = fileObj.inode.mtime()

            flavor -= fileObj.flavor()
            self._componentReqs[troveName] -= fileObj.requires()
            self._componentProvs[troveName] -= fileObj.requires()

            # Config vs. InitialContents etc. might be change in derived pkg
            # Set defaults here, and they can be overridden with
            # "exceptions = " later
            if fileObj.flags.isConfig():
                self.Config(path)
            elif fileObj.flags.isInitialContents():
                self.InitialContents(path)
            elif fileObj.flags.isTransient():
                self.Transient(path)

            # we don't restore setuid/setgid bits into the filesystem
            if fileObj.inode.perms() & 06000 != 0:
                self.SetModes(path, fileObj.inode.perms())

            if isinstance(fileObj, files.DeviceFile):
                self.MakeDevices(path, fileObj.lsTag,
                                 fileObj.devt.major(), fileObj.devt.minor(),
                                 fileObj.inode.owner(), fileObj.inode.group(),
                                 fileObj.inode.perms())
            else:
                if fileObj.hasContents:
                    linkGroup = fileObj.linkGroup()
                    if linkGroup:
                        l = linkGroups.setdefault(linkGroup, [])
                        l.append(path)

                    (contentType, contents) = \
                                    self.cs.getFileContents(pathId, fileId)
                    if contentType == changeset.ChangedFileTypes.ptr:
                        targetPtrId = contents.get().read()
                        l = delayedRestores.setdefault(targetPtrId, [])
                        l.append((fileObj, path))
                        continue
                    elif linkGroup and not linkGroup in linkGroupFirstPath:
                        # only non-delayed restores can be initial target
                        linkGroupFirstPath[linkGroup] = path

                    assert(contentType == changeset.ChangedFileTypes.file)
                else:
                    contents = None

                ptrId = pathId + fileId
                if pathId in delayedRestores:
                    ptrMap[pathId] = path
                elif ptrId in delayedRestores:
                    ptrMap[ptrId] = path

                fileObj.restore(contents, destdir, destdir + path)

            if isinstance(fileObj, files.Directory):
                # remember to include this directory in the derived package
                self.ExcludeDirectories(exceptions = path)

        for targetPtrId in delayedRestores:
            for fileObj, targetPath in delayedRestores[targetPtrId]:
                sourcePath = ptrMap[targetPtrId]
                fileObj.restore(
                    filecontents.FromFilesystem(destdir + sourcePath),
                    destdir, destdir + targetPath)

        # we do not have to worry about cross-device hardlinks in destdir
        for linkGroup in linkGroups:
            for path in linkGroups[linkGroup]:
                initialPath = linkGroupFirstPath[linkGroup]
                if path == initialPath:
                    continue
                util.createLink(destdir + initialPath, destdir + path)

        self.useFlags = flavor

        self.setByDefaultOn(set(x for x in byDefault if byDefault[x]))
        self.setByDefaultOff(set(x for x in byDefault if not byDefault[x]))

    def unpackSources(self, builddir, destdir, resume=None,
                      downloadOnly=False):
        if self.parentVersion:
            try:
                parentRevision = versions.Revision(self.parentVersion)
            except conaryerrors.ParseError, e:
                raise builderrors.RecipeFileError(
                            'Cannot parse parentVersion %s: %s',
                                    self.parentRevision, str(e))
        else:
            parentRevision = None

        if not self.sourceVersion.isShadow():
            raise builderrors.RecipeFileError(
                    "only shadowed sources can be derived packages")

        if parentRevision and \
                self.sourceVersion.trailingRevision().getVersion() != \
                                                parentRevision.getVersion():
            raise builderrors.RecipeFileError(
                    "parentRevision must have the same upstream version as the "
                    "derived package recipe")

        # find all the flavors of the parent
        parentBranch = self.sourceVersion.branch().parentBranch()

        if parentRevision:
            parentVersion = parentBranch.createVersion(parentRevision)

            d = self.repos.getTroveVersionFlavors({ self.name :
                                { parentVersion : [ None ] } } )
            if self.name not in d:
                raise builderrors.RecipeFileError(
                        'Version %s of %s not found'
                                    % (parentVersion, self.name) )
        else:
            d = self.repos.getTroveLeavesByBranch(
                    { self.name : { parentBranch : [ None ] } } )

            if not d[self.name]:
                raise builderrors.RecipeFileError(
                    'No versions of %s found on branch %s' %
                            (self.name, parentBranch))

            parentVersion = sorted(d[self.name].keys())[-1]

        bestFlavor = (-1, [])
        # choose which flavor to derive from
        for flavor in d[self.name][parentVersion]:
            score = self.cfg.buildFlavor.score(flavor)
            if score is False: continue
            if bestFlavor[0] < score:
                bestFlavor = (score, [ flavor ])
            elif bestFlavor[0] == score:
                bestFlavor[1].append(flavor)

        if bestFlavor[0] == -1:
            raise builderrors.RecipeFileError(
                    'No flavors of %s=%s found for build flavor %s',
                    self.name, parentVersion, self.cfg.buildFlavor)
        elif len(bestFlavor[1]) > 1:
            raise builderrors.RecipeFileError(
                    'Multiple flavors of %s=%s match build flavor %s',
                    self.name, parentVersion, self.cfg.buildFlavor)

        parentFlavor = bestFlavor[1][0]

        log.info('deriving from %s=%s[%s]', self.name, parentVersion,
                 parentFlavor)

        # Fetch all binaries built from this source
        binaries = self.repos.getTrovesBySource(self.name + ':source',
                parentVersion.getSourceVersion())

        # Filter out older ones
        binaries = [ x for x in binaries if x[1] == parentVersion ]

        # Build trove spec
        troveSpec = [ (x[0], (None, None), (x[1], x[2]), True)
                        for x in binaries ]

        self.cs = self.repos.createChangeSet(troveSpec, recurse = False)
        self.addLoadedTroves([
            (x.getName(), x.getNewVersion(), x.getNewFlavor()) for x
            in self.cs.iterNewTroveList() ])

        self._expandChangeset()

        _AbstractPackageRecipe.unpackSources(self, builddir, destdir,
                                             resume = resume,
                                             downloadOnly = downloadOnly)

    def loadPolicy(self):
        return _AbstractPackageRecipe.loadPolicy(self,
                                internalPolicyModules = ( 'derivedpolicy', ) )

    def __init__(self, cfg, laReposCache, srcDirs, extraMacros={},
                 crossCompile=None, lightInstance=False):
        _AbstractPackageRecipe.__init__(self, cfg, laReposCache, srcDirs,
                                        extraMacros = extraMacros,
                                        crossCompile = crossCompile,
                                        lightInstance = lightInstance)

        log.info('Warning: Derived packages are experimental and subject to change')

        self.repos = laReposCache.repos

        self._addBuildAction('Ant', build.Ant)
        self._addBuildAction('Automake', build.Automake)
        self._addBuildAction('ClassPath', build.ClassPath)
        self._addBuildAction('CompilePython', build.CompilePython)
        self._addBuildAction('Configure', build.Configure)
        self._addBuildAction('ConsoleHelper', build.ConsoleHelper)
        self._addBuildAction('Copy', build.Copy)
        self._addBuildAction('Create', build.Create)
        self._addBuildAction('Desktopfile', build.Desktopfile)
        self._addBuildAction('Doc', build.Doc)
        self._addBuildAction('Environment', build.Environment)
        self._addBuildAction('Install', build.Install)
        self._addBuildAction('JavaCompile', build.JavaCompile)
        self._addBuildAction('JavaDoc', build.JavaDoc)
        self._addBuildAction('Link', build.Link)
        self._addBuildAction('Make', build.Make)
        self._addBuildAction('MakeDirs', build.MakeDirs)
        self._addBuildAction('MakeInstall', build.MakeInstall)
        self._addBuildAction('MakeParallelSubdir', build.MakeParallelSubdir)
        self._addBuildAction('MakePathsInstall', build.MakePathsInstall)
        self._addBuildAction('ManualConfigure', build.ManualConfigure)
        self._addBuildAction('Move', build.Move)
        self._addBuildAction('PythonSetup', build.PythonSetup)
        self._addBuildAction('Remove', build.Remove)
        self._addBuildAction('Replace', build.Replace)
        self._addBuildAction('Run', build.Run)
        self._addBuildAction('SetModes', build.SetModes)
        self._addBuildAction('SGMLCatalogEntry', build.SGMLCatalogEntry)
        self._addBuildAction('Symlink', build.Symlink)
        self._addBuildAction('XInetdService', build.XInetdService)
        self._addBuildAction('XMLCatalogEntry', build.XMLCatalogEntry)

        self._addSourceAction('addArchive', source.addArchive)
        self._addSourceAction('addAction', source.addAction)
        self._addSourceAction('addPatch', source.addPatch)
        self._addSourceAction('addSource', source.addSource)
<commit_msg>Filtering by parent flavor (CNY-1423)<commit_after># Copyright (c) 2006-2007 rPath, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.opensource.org/licenses/cpl.php.
#
# This program 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 Common Public License for
# full details.

from conary import files, trove, versions
from conary import errors as conaryerrors
from conary.build import build, source
from conary.build import errors as builderrors
from conary.build.packagerecipe import _AbstractPackageRecipe
from conary.local import update
from conary.lib import log, util
from conary.repository import changeset, filecontents

class DerivedPackageRecipe(_AbstractPackageRecipe):

    internalAbstractBaseClass = 1
    _isDerived = True
    parentVersion = None

    def _expandChangeset(self):
        destdir = self.macros.destdir

        delayedRestores = {}
        ptrMap = {}
        byDefault = {}

        fileList = []
        linkGroups = {}
        linkGroupFirstPath = {}
        # sort the files by pathId,fileId
        for trvCs in self.cs.iterNewTroveList():
            trv = trove.Trove(trvCs)

            # these should all be the same anyway
            flavor = trv.getFlavor().copy()
            name = trv.getName()
            self._componentReqs[name] = trv.getRequires().copy()
            self._componentProvs[name] = trv.getProvides().copy()

            for pathId, path, fileId, version in trv.iterFileList():
                if path != self.macros.buildlogpath:
                    fileList.append((pathId, fileId, path, name))

            if trv.isCollection():
                # gather up existing byDefault status
                # from (component, byDefault) tuples
                byDefault.update(dict(
                    [(x[0][0], x[1]) for x in trv.iterTroveListInfo()]))

        fileList.sort()

        for pathId, fileId, path, troveName in fileList:
            fileCs = self.cs.getFileChange(None, fileId)
            fileObj = files.ThawFile(fileCs, pathId)
            self._derivedFiles[path] = fileObj.inode.mtime()

            flavor -= fileObj.flavor()
            self._componentReqs[troveName] -= fileObj.requires()
            self._componentProvs[troveName] -= fileObj.requires()

            # Config vs. InitialContents etc. might be change in derived pkg
            # Set defaults here, and they can be overridden with
            # "exceptions = " later
            if fileObj.flags.isConfig():
                self.Config(path)
            elif fileObj.flags.isInitialContents():
                self.InitialContents(path)
            elif fileObj.flags.isTransient():
                self.Transient(path)

            # we don't restore setuid/setgid bits into the filesystem
            if fileObj.inode.perms() & 06000 != 0:
                self.SetModes(path, fileObj.inode.perms())

            if isinstance(fileObj, files.DeviceFile):
                self.MakeDevices(path, fileObj.lsTag,
                                 fileObj.devt.major(), fileObj.devt.minor(),
                                 fileObj.inode.owner(), fileObj.inode.group(),
                                 fileObj.inode.perms())
            else:
                if fileObj.hasContents:
                    linkGroup = fileObj.linkGroup()
                    if linkGroup:
                        l = linkGroups.setdefault(linkGroup, [])
                        l.append(path)

                    (contentType, contents) = \
                                    self.cs.getFileContents(pathId, fileId)
                    if contentType == changeset.ChangedFileTypes.ptr:
                        targetPtrId = contents.get().read()
                        l = delayedRestores.setdefault(targetPtrId, [])
                        l.append((fileObj, path))
                        continue
                    elif linkGroup and not linkGroup in linkGroupFirstPath:
                        # only non-delayed restores can be initial target
                        linkGroupFirstPath[linkGroup] = path

                    assert(contentType == changeset.ChangedFileTypes.file)
                else:
                    contents = None

                ptrId = pathId + fileId
                if pathId in delayedRestores:
                    ptrMap[pathId] = path
                elif ptrId in delayedRestores:
                    ptrMap[ptrId] = path

                fileObj.restore(contents, destdir, destdir + path)

            if isinstance(fileObj, files.Directory):
                # remember to include this directory in the derived package
                self.ExcludeDirectories(exceptions = path)

        for targetPtrId in delayedRestores:
            for fileObj, targetPath in delayedRestores[targetPtrId]:
                sourcePath = ptrMap[targetPtrId]
                fileObj.restore(
                    filecontents.FromFilesystem(destdir + sourcePath),
                    destdir, destdir + targetPath)

        # we do not have to worry about cross-device hardlinks in destdir
        for linkGroup in linkGroups:
            for path in linkGroups[linkGroup]:
                initialPath = linkGroupFirstPath[linkGroup]
                if path == initialPath:
                    continue
                util.createLink(destdir + initialPath, destdir + path)

        self.useFlags = flavor

        self.setByDefaultOn(set(x for x in byDefault if byDefault[x]))
        self.setByDefaultOff(set(x for x in byDefault if not byDefault[x]))

    def unpackSources(self, builddir, destdir, resume=None,
                      downloadOnly=False):
        if self.parentVersion:
            try:
                parentRevision = versions.Revision(self.parentVersion)
            except conaryerrors.ParseError, e:
                raise builderrors.RecipeFileError(
                            'Cannot parse parentVersion %s: %s',
                                    self.parentRevision, str(e))
        else:
            parentRevision = None

        if not self.sourceVersion.isShadow():
            raise builderrors.RecipeFileError(
                    "only shadowed sources can be derived packages")

        if parentRevision and \
                self.sourceVersion.trailingRevision().getVersion() != \
                                                parentRevision.getVersion():
            raise builderrors.RecipeFileError(
                    "parentRevision must have the same upstream version as the "
                    "derived package recipe")

        # find all the flavors of the parent
        parentBranch = self.sourceVersion.branch().parentBranch()

        if parentRevision:
            parentVersion = parentBranch.createVersion(parentRevision)

            d = self.repos.getTroveVersionFlavors({ self.name :
                                { parentVersion : [ None ] } } )
            if self.name not in d:
                raise builderrors.RecipeFileError(
                        'Version %s of %s not found'
                                    % (parentVersion, self.name) )
        else:
            d = self.repos.getTroveLeavesByBranch(
                    { self.name : { parentBranch : [ None ] } } )

            if not d[self.name]:
                raise builderrors.RecipeFileError(
                    'No versions of %s found on branch %s' %
                            (self.name, parentBranch))

            parentVersion = sorted(d[self.name].keys())[-1]

        bestFlavor = (-1, [])
        # choose which flavor to derive from
        for flavor in d[self.name][parentVersion]:
            score = self.cfg.buildFlavor.score(flavor)
            if score is False: continue
            if bestFlavor[0] < score:
                bestFlavor = (score, [ flavor ])
            elif bestFlavor[0] == score:
                bestFlavor[1].append(flavor)

        if bestFlavor[0] == -1:
            raise builderrors.RecipeFileError(
                    'No flavors of %s=%s found for build flavor %s',
                    self.name, parentVersion, self.cfg.buildFlavor)
        elif len(bestFlavor[1]) > 1:
            raise builderrors.RecipeFileError(
                    'Multiple flavors of %s=%s match build flavor %s',
                    self.name, parentVersion, self.cfg.buildFlavor)

        parentFlavor = bestFlavor[1][0]

        log.info('deriving from %s=%s[%s]', self.name, parentVersion,
                 parentFlavor)

        # Fetch all binaries built from this source
        binaries = self.repos.getTrovesBySource(self.name + ':source',
                parentVersion.getSourceVersion())

        # Filter out older ones
        binaries = [ x for x in binaries \
                        if (x[1], x[2]) == (parentVersion, parentFlavor) ]

        # Build trove spec
        troveSpec = [ (x[0], (None, None), (x[1], x[2]), True)
                        for x in binaries ]

        self.cs = self.repos.createChangeSet(troveSpec, recurse = False)
        self.addLoadedTroves([
            (x.getName(), x.getNewVersion(), x.getNewFlavor()) for x
            in self.cs.iterNewTroveList() ])

        self._expandChangeset()

        _AbstractPackageRecipe.unpackSources(self, builddir, destdir,
                                             resume = resume,
                                             downloadOnly = downloadOnly)

    def loadPolicy(self):
        return _AbstractPackageRecipe.loadPolicy(self,
                                internalPolicyModules = ( 'derivedpolicy', ) )

    def __init__(self, cfg, laReposCache, srcDirs, extraMacros={},
                 crossCompile=None, lightInstance=False):
        _AbstractPackageRecipe.__init__(self, cfg, laReposCache, srcDirs,
                                        extraMacros = extraMacros,
                                        crossCompile = crossCompile,
                                        lightInstance = lightInstance)

        log.info('Warning: Derived packages are experimental and subject to change')

        self.repos = laReposCache.repos

        self._addBuildAction('Ant', build.Ant)
        self._addBuildAction('Automake', build.Automake)
        self._addBuildAction('ClassPath', build.ClassPath)
        self._addBuildAction('CompilePython', build.CompilePython)
        self._addBuildAction('Configure', build.Configure)
        self._addBuildAction('ConsoleHelper', build.ConsoleHelper)
        self._addBuildAction('Copy', build.Copy)
        self._addBuildAction('Create', build.Create)
        self._addBuildAction('Desktopfile', build.Desktopfile)
        self._addBuildAction('Doc', build.Doc)
        self._addBuildAction('Environment', build.Environment)
        self._addBuildAction('Install', build.Install)
        self._addBuildAction('JavaCompile', build.JavaCompile)
        self._addBuildAction('JavaDoc', build.JavaDoc)
        self._addBuildAction('Link', build.Link)
        self._addBuildAction('Make', build.Make)
        self._addBuildAction('MakeDirs', build.MakeDirs)
        self._addBuildAction('MakeInstall', build.MakeInstall)
        self._addBuildAction('MakeParallelSubdir', build.MakeParallelSubdir)
        self._addBuildAction('MakePathsInstall', build.MakePathsInstall)
        self._addBuildAction('ManualConfigure', build.ManualConfigure)
        self._addBuildAction('Move', build.Move)
        self._addBuildAction('PythonSetup', build.PythonSetup)
        self._addBuildAction('Remove', build.Remove)
        self._addBuildAction('Replace', build.Replace)
        self._addBuildAction('Run', build.Run)
        self._addBuildAction('SetModes', build.SetModes)
        self._addBuildAction('SGMLCatalogEntry', build.SGMLCatalogEntry)
        self._addBuildAction('Symlink', build.Symlink)
        self._addBuildAction('XInetdService', build.XInetdService)
        self._addBuildAction('XMLCatalogEntry', build.XMLCatalogEntry)

        self._addSourceAction('addArchive', source.addArchive)
        self._addSourceAction('addAction', source.addAction)
        self._addSourceAction('addPatch', source.addPatch)
        self._addSourceAction('addSource', source.addSource)
<|endoftext|>