text:
import arrow
import uuid
import re
from datetime import datetime, timezone
from astropy import units as u
from astropy import time as ap_time
import astroplan
import numpy as np
import sqlalchemy as sa
from sqlalchemy import cast
from sqlalchemy.dialects.postgresql import ARRAY
from sqlalchemy.dialects import postgresql as psql
from sqlalchemy.orm import relationship
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy_utils import ArrowType, URLType

from astropy import coordinates as ap_coord
import healpix_alchemy as ha

from baselayer.app.models import (  # noqa
    init_db,
    join_model,
    Base,
    DBSession,
    ACL,
    Role,
    User,
    Token,
)
from baselayer.app.custom_exceptions import AccessError

import timezonefinder

from . import schema
from .enum_types import (
    allowed_bandpasses,
    thumbnail_types,
    instrument_types,
    followup_priorities,
)

# In the AB system, a brightness of 23.9 mag corresponds to 1 microJy.
# All DB fluxes are stored in microJy (AB).
PHOT_ZP = 23.9
PHOT_SYS = 'ab'

def is_owned_by(self, user_or_token):
    """Generic ownership logic for any `skyportal` ORM model.

    Models with complicated ownership logic should implement their own method
    instead of adding too many additional conditions here.
    """
    if hasattr(self, 'tokens'):
        return user_or_token in self.tokens
    if hasattr(self, 'groups'):
        return bool(set(self.groups) & set(user_or_token.accessible_groups))
    if hasattr(self, 'group'):
        return self.group in user_or_token.accessible_groups
    if hasattr(self, 'users'):
        if hasattr(user_or_token, 'created_by'):
            if user_or_token.created_by in self.users:
                return True
        return user_or_token in self.users

    raise NotImplementedError(f"{type(self).__name__} object has no owner")

Base.is_owned_by = is_owned_by

class NumpyArray(sa.types.TypeDecorator):
    impl = psql.ARRAY(sa.Float)

    def process_result_value(self, value, dialect):
        return np.array(value)

class Group(Base):
    name = sa.Column(sa.String, unique=True, nullable=False)

    streams = relationship(
        'Stream',
        secondary='group_streams',
        back_populates='groups',
        passive_deletes=True,
    )
    filters = relationship("Filter", back_populates="group", passive_deletes=True)

    users = relationship(
        'User', secondary='group_users', back_populates='groups', passive_deletes=True
    )
    group_users = relationship(
        'GroupUser',
        back_populates='group',
        cascade='save-update, merge, refresh-expire, expunge',
        passive_deletes=True,
    )

    observing_runs = relationship('ObservingRun', back_populates='group')
    photometry = relationship(
        "Photometry",
        secondary="group_photometry",
        back_populates="groups",
        cascade="save-update, merge, refresh-expire, expunge",
        passive_deletes=True,
    )
    spectra = relationship(
        "Spectrum",
        secondary="group_spectra",
        back_populates="groups",
        cascade="save-update, merge, refresh-expire, expunge",
        passive_deletes=True,
    )
    single_user_group = sa.Column(sa.Boolean, default=False)

GroupUser = join_model('group_users', Group, User)
GroupUser.admin = sa.Column(sa.Boolean, nullable=False, default=False)

class Stream(Base):
    name = sa.Column(sa.String, unique=True, nullable=False)
    altdata = sa.Column(
        JSONB,
        nullable=True,
        doc="Misc. metadata stored in JSON format, e.g. "
        "`{'collection': 'ZTF_alerts', selector: [1, 2]}`",
    )

    groups = relationship(
        'Group',
        secondary='group_streams',
        back_populates='streams',
        passive_deletes=True,
    )
    users = relationship(
        'User', secondary='stream_users', back_populates='streams', passive_deletes=True
    )
    filters = relationship('Filter', back_populates='stream', passive_deletes=True)

GroupStream = join_model('group_streams', Group, Stream)

StreamUser = join_model('stream_users', Stream, User)

User.groups = relationship(
    'Group', secondary='group_users', back_populates='users', passive_deletes=True
)

User.streams = relationship(
    'Stream', secondary='stream_users', back_populates='users', passive_deletes=True
)

@property
def user_or_token_accessible_groups(self):
    if "System admin" in [acl.id for acl in self.acls]:
        return Group.query.all()
    return self.groups

User.accessible_groups = user_or_token_accessible_groups
Token.accessible_groups = user_or_token_accessible_groups

@property
def token_groups(self):
    return self.created_by.groups

Token.groups = token_groups

class Obj(Base, ha.Point):
    id = sa.Column(sa.String, primary_key=True)
    # TODO should this column type be decimal? fixed-precison numeric

    ra_dis = sa.Column(sa.Float)
    dec_dis = sa.Column(sa.Float)

    ra_err = sa.Column(sa.Float, nullable=True)
    dec_err = sa.Column(sa.Float, nullable=True)

    offset = sa.Column(sa.Float, default=0.0)
    redshift = sa.Column(sa.Float, nullable=True)

    # Contains all external metadata, e.g. simbad, pan-starrs, tns, gaia
    altdata = sa.Column(
        JSONB,
        nullable=True,
        doc="Misc. alternative metadata stored in JSON format, e.g. "
        "`{'gaia': {'info': {'Teff': 5780}}}`",
    )

    dist_nearest_source = sa.Column(sa.Float, nullable=True)
    mag_nearest_source = sa.Column(sa.Float, nullable=True)
    e_mag_nearest_source = sa.Column(sa.Float, nullable=True)

    transient = sa.Column(sa.Boolean, default=False)
    varstar = sa.Column(sa.Boolean, default=False)
    is_roid = sa.Column(sa.Boolean, default=False)

    score = sa.Column(sa.Float, nullable=True)

    origin = sa.Column(sa.String, nullable=True)

    comments = relationship(
        'Comment',
        back_populates='obj',
        cascade='save-update, merge, refresh-expire, expunge',
        passive_deletes=True,
        order_by="Comment.created_at",
    )

    classifications = relationship(
        'Classification',
        back_populates='obj',
        cascade='save-update, merge, refresh-expire, expunge',
        passive_deletes=True,
        order_by="Classification.created_at",
    )

    photometry = relationship(
        'Photometry',
        back_populates='obj',
        cascade='save-update, merge, refresh-expire, expunge',
        single_parent=True,
        passive_deletes=True,
        order_by="Photometry.mjd",
    )

    detect_photometry_count = sa.Column(sa.Integer, nullable=True)

    spectra = relationship(
        'Spectrum',
        back_populates='obj',
        cascade='save-update, merge, refresh-expire, expunge',
        single_parent=True,
        passive_deletes=True,
        order_by="Spectrum.observed_at",
    )
    thumbnails = relationship(
        'Thumbnail',
        back_populates='obj',
        secondary='photometry',
        cascade='save-update, merge, refresh-expire, expunge',
        passive_deletes=True,
    )

    followup_requests = relationship('FollowupRequest', back_populates='obj')
    assignments = relationship('ClassicalAssignment', back_populates='obj')

    @hybrid_property
    def last_detected(self):
        detections = [phot.iso for phot in self.photometry if phot.snr and phot.snr > 5]
        return max(detections) if detections else None

    @last_detected.expression
    def last_detected(cls):
        return (
            sa.select([sa.func.max(Photometry.iso)])
            .where(Photometry.obj_id == cls.id)
            .where(Photometry.snr > 5.0)
            .group_by(Photometry.obj_id)
            .label('last_detected')
        )

    def add_linked_thumbnails(self):
        sdss_thumb = Thumbnail(
            photometry=self.photometry[0], public_url=self.sdss_url, type='sdss'
        )
        dr8_thumb = Thumbnail(
            photometry=self.photometry[0], public_url=self.desi_dr8_url, type='dr8'
        )
        DBSession().add_all([sdss_thumb, dr8_thumb])
        DBSession().commit()

    @property
    def sdss_url(self):
        """Construct URL for public Sloan Digital Sky Survey (SDSS) cutout."""
        return (
            f"http://skyserver.sdss.org/dr12/SkyserverWS/ImgCutout/getjpeg"
            f"?ra={self.ra}&dec={self.dec}&scale=0.3&width=200&height=200"
            f"&opt=G&query=&Grid=on"
        )

    @property
    def desi_dr8_url(self):
        """Construct URL for public DESI DR8 cutout."""
        return (
            f"http://legacysurvey.org/viewer/jpeg-cutout?ra={self.ra}"
            f"&dec={self.dec}&size=200&layer=dr8&pixscale=0.262&bands=grz"
        )

    @property
    def target(self):
        """Representation of this Obj as an astroplan.FixedTarget."""
        coord = ap_coord.SkyCoord(self.ra, self.dec, unit='deg')
        return astroplan.FixedTarget(name=self.id, coord=coord)

    @property
    def gal_lat_deg(self):
        """Get the galactic latitute of this object"""
        coord = ap_coord.SkyCoord(self.ra, self.dec, unit="deg")
        return coord.galactic.b.deg

    @property
    def gal_lon_deg(self):
        """Get the galactic longitude of this object"""
        coord = ap_coord.SkyCoord(self.ra, self.dec, unit="deg")
        return coord.galactic.l.deg

    def airmass(self, telescope, time, below_horizon=np.inf):
        """Return the airmass of the object at a given time. Uses the Pickering
        (2002) interpolation of the Rayleigh (molecular atmosphere) airmass.

        The Pickering interpolation tends toward 38.7494 as the altitude
        approaches zero.

        Parameters
        ----------
        telescope : `skyportal.models.Telescope`
            The telescope to use for the airmass calculation
        time : `astropy.time.Time` or list of astropy.time.Time`
            The time or times at which to calculate the airmass
        below_horizon : scalar, Numeric
            Airmass value to assign when an object is below the horizon.
            An object is "below the horizon" when its altitude is less than
            zero degrees.

        Returns
        -------
        airmass : ndarray
           The airmass of the Obj at the requested times
        """

        output_shape = np.shape(time)
        time = np.atleast_1d(time)
        altitude = self.altitude(telescope, time).to('degree').value
        above = altitude > 0

        # use Pickering (2002) interpolation to calculate the airmass
        # The Pickering interpolation tends toward 38.7494 as the altitude
        # approaches zero.
        sinarg = np.zeros_like(altitude)
        airmass = np.ones_like(altitude) * np.inf
        sinarg[above] = altitude[above] + 244 / (165 + 47 * altitude[above] ** 1.1)
        airmass[above] = 1.0 / np.sin(np.deg2rad(sinarg[above]))

        # set objects below the horizon to an airmass of infinity
        airmass[~above] = below_horizon
        airmass = airmass.reshape(output_shape)

        return airmass

    def altitude(self, telescope, time):
        """Return the altitude of the object at a given time.

        Parameters
        ----------
        telescope : `skyportal.models.Telescope`
            The telescope to use for the altitude calculation

        time : `astropy.time.Time`
            The time or times at which to calculate the altitude

        Returns
        -------
        alt : `astropy.coordinates.AltAz`
           The altitude of the Obj at the requested times
        """

        return telescope.observer.altaz(time, self.target).alt

class Filter(Base):
    name = sa.Column(sa.String, nullable=False, unique=False)
    stream_id = sa.Column(
        sa.ForeignKey("streams.id", ondelete="CASCADE"), nullable=False, index=True
    )
    stream = relationship("Stream", foreign_keys=[stream_id], back_populates="filters")
    group_id = sa.Column(
        sa.ForeignKey("groups.id", ondelete="CASCADE"), nullable=False, index=True
    )
    group = relationship("Group", foreign_keys=[group_id], back_populates="filters")

Candidate = join_model("candidates", Filter, Obj)
Candidate.passed_at = sa.Column(sa.DateTime, nullable=True)
Candidate.passing_alert_id = sa.Column(sa.BigInteger)

def get_candidate_if_owned_by(obj_id, user_or_token, options=[]):
    if Candidate.query.filter(Candidate.obj_id == obj_id).first() is None:
        return None
    user_group_ids = [g.id for g in user_or_token.accessible_groups]
    c = (
        Candidate.query.filter(Candidate.obj_id == obj_id)
        .filter(
            Candidate.filter_id.in_(
                DBSession.query(Filter.id).filter(Filter.group_id.in_(user_group_ids))
            )
        )
        .options(options)
        .first()
    )
    if c is None:
        raise AccessError("Insufficient permissions.")
    return c.obj

def candidate_is_owned_by(self, user_or_token):
    return self.filter.group in user_or_token.accessible_groups

Candidate.get_obj_if_owned_by = get_candidate_if_owned_by
Candidate.is_owned_by = candidate_is_owned_by

Source = join_model("sources", Group, Obj)
"""User.sources defines the logic for whether a user has access to a source;
   if this gets more complicated it should become a function/`hybrid_property`
   rather than a `relationship`.
"""
Source.saved_by_id = sa.Column(
    sa.ForeignKey("users.id"), nullable=True, unique=False, index=True
)
Source.saved_by = relationship(
    "User", foreign_keys=[Source.saved_by_id], backref="saved_sources"
)
Source.saved_at = sa.Column(sa.DateTime, nullable=True)
Source.active = sa.Column(sa.Boolean, server_default="true")
Source.requested = sa.Column(sa.Boolean, server_default="false")
Source.unsaved_by_id = sa.Column(
    sa.ForeignKey("users.id"), nullable=True, unique=False, index=True
)
Source.unsaved_by = relationship("User", foreign_keys=[Source.unsaved_by_id])

def source_is_owned_by(self, user_or_token):
    source_group_ids = [
        row[0]
        for row in DBSession.query(Source.group_id)
        .filter(Source.obj_id == self.obj_id)
        .all()
    ]
    return bool(set(source_group_ids) & {g.id for g in user_or_token.accessible_groups})

def get_source_if_owned_by(obj_id, user_or_token, options=[]):
    if Source.query.filter(Source.obj_id == obj_id).first() is None:
        return None
    user_group_ids = [g.id for g in user_or_token.accessible_groups]
    s = (
        Source.query.filter(Source.obj_id == obj_id)
        .filter(Source.group_id.in_(user_group_ids))
        .options(options)
        .first()
    )
    if s is None:
        raise AccessError("Insufficient permissions.")
    return s.obj

Source.is_owned_by = source_is_owned_by
Source.get_obj_if_owned_by = get_source_if_owned_by

def get_obj_if_owned_by(obj_id, user_or_token, options=[]):
    if Obj.query.get(obj_id) is None:
        return None
    try:
        obj = Source.get_obj_if_owned_by(obj_id, user_or_token, options)
    except AccessError:  # They may still be able to view the associated Candidate
        obj = Candidate.get_obj_if_owned_by(obj_id, user_or_token, options)
        if obj is None:
            # If user can't view associated Source, and there's no Candidate they can
            # view, raise AccessError
            raise
    if obj is None:  # There is no associated Source/Cand, so check based on photometry
        if Obj.get_photometry_owned_by_user(obj_id, user_or_token):
            return Obj.query.options(options).get(obj_id)
        raise AccessError("Insufficient permissions.")
    # If we get here, the user has access to either the associated Source or Candidate
    return obj

Obj.get_if_owned_by = get_obj_if_owned_by

def get_obj_comments_owned_by(self, user_or_token):
    owned_comments = [
        comment for comment in self.comments if comment.is_owned_by(user_or_token)
    ]

    # Grab basic author info for the comments
    for comment in owned_comments:
        comment.author_info = comment.construct_author_info_dict()

    return owned_comments

Obj.get_comments_owned_by = get_obj_comments_owned_by

def get_obj_classifications_owned_by(self, user_or_token):
    return [
        classifications
        for classifications in self.classifications
        if classifications.is_owned_by(user_or_token)
    ]

Obj.get_classifications_owned_by = get_obj_classifications_owned_by

def get_photometry_owned_by_user(obj_id, user_or_token):
    return (
        Photometry.query.filter(Photometry.obj_id == obj_id)
        .filter(
            Photometry.groups.any(
                Group.id.in_([g.id for g in user_or_token.accessible_groups])
            )
        )
        .all()
    )

Obj.get_photometry_owned_by_user = get_photometry_owned_by_user

def get_spectra_owned_by(obj_id, user_or_token):
    return (
        Spectrum.query.filter(Spectrum.obj_id == obj_id)
        .filter(
            Spectrum.groups.any(
                Group.id.in_([g.id for g in user_or_token.accessible_groups])
            )
        )
        .all()
    )

Obj.get_spectra_owned_by = get_spectra_owned_by

User.sources = relationship(
    'Obj',
    backref='users',
    secondary='join(Group, sources).join(group_users)',
    primaryjoin='group_users.c.user_id == users.c.id',
    passive_deletes=True,
)

class SourceView(Base):
    obj_id = sa.Column(
        sa.ForeignKey('objs.id', ondelete='CASCADE'),
        nullable=False,
        unique=False,
        index=True,
    )
    username_or_token_id = sa.Column(sa.String, nullable=False, unique=False)
    is_token = sa.Column(sa.Boolean, nullable=False, default=False)
    created_at = sa.Column(
        sa.DateTime, nullable=False, default=datetime.now, index=True
    )

class Telescope(Base):
    name = sa.Column(sa.String, unique=True, nullable=False)
    nickname = sa.Column(sa.String, nullable=False)
    lat = sa.Column(sa.Float, nullable=True, doc='Latitude in deg.')
    lon = sa.Column(sa.Float, nullable=True, doc='Longitude in deg.')
    elevation = sa.Column(sa.Float, nullable=True, doc='Elevation in meters.')
    diameter = sa.Column(sa.Float, nullable=False, doc='Diameter in meters.')
    skycam_link = sa.Column(
        URLType, nullable=True, doc="Link to the telescope's sky camera."
    )
    robotic = sa.Column(
        sa.Boolean, default=False, nullable=False, doc="Is this telescope robotic?"
    )

    instruments = relationship(
        'Instrument',
        back_populates='telescope',
        cascade='save-update, merge, refresh-expire, expunge',
        passive_deletes=True,
    )

    @property
    def observer(self):
        tf = timezonefinder.TimezoneFinder()
        local_tz = tf.timezone_at(lng=self.lon, lat=self.lat)
        return astroplan.Observer(
            longitude=self.lon * u.deg,
            latitude=self.lat * u.deg,
            elevation=self.elevation * u.m,
            timezone=local_tz,
        )

class ArrayOfEnum(ARRAY):
    def bind_expression(self, bindvalue):
        return cast(bindvalue, self)

    def result_processor(self, dialect, coltype):
        super_rp = super(ArrayOfEnum, self).result_processor(dialect, coltype)

        def handle_raw_string(value):
            if value is None or value == '{}':  # 2nd case, empty array
                return []
            inner = re.match(r"^{(.*)}$", value).group(1)
            return inner.split(",")

        def process(value):
            return super_rp(handle_raw_string(value))

        return process

class Instrument(Base):
    name = sa.Column(sa.String, unique=True, nullable=False)
    type = sa.Column(instrument_types, nullable=False)

    band = sa.Column(sa.String)
    telescope_id = sa.Column(
        sa.ForeignKey('telescopes.id', ondelete='CASCADE'), nullable=False, index=True
    )
    telescope = relationship('Telescope', back_populates='instruments')

    followup_requests = relationship('FollowupRequest', back_populates='instrument')

    photometry = relationship('Photometry', back_populates='instrument')
    spectra = relationship('Spectrum', back_populates='instrument')

    # can be [] if an instrument is spec only
    filters = sa.Column(
        ArrayOfEnum(allowed_bandpasses),
        nullable=False,
        default=[],
        doc='List of filters on the instrument ' '(if any).',
    )

    observing_runs = relationship('ObservingRun', back_populates='instrument')

    @property
    def does_spectroscopy(self):
        return 'spec' in self.type

    @property
    def does_imaging(self):
        return 'imag' in self.type

class Taxonomy(Base):
    __tablename__ = 'taxonomies'
    name = sa.Column(
        sa.String,
        nullable=False,
        doc='Short string to make this taxonomy memorable ' 'to end users.',
    )
    hierarchy = sa.Column(
        JSONB,
        nullable=False,
        doc='Nested JSON describing the taxonomy '
        'which should be validated against '
        'a schema before entry',
    )
    provenance = sa.Column(
        sa.String,
        nullable=True,
        doc='Identifier (e.g., URL or git hash) that '
        'uniquely ties this taxonomy back '
        'to an origin or place of record',
    )
    version = sa.Column(
        sa.String, nullable=False, doc='Semantic version of this taxonomy'
    )

    isLatest = sa.Column(
        sa.Boolean,
        default=True,
        nullable=False,
        doc='Consider this the latest version of '
        'the taxonomy with this name? Defaults '
        'to True.',
    )
    groups = relationship(
        "Group",
        secondary="group_taxonomy",
        cascade="save-update," "merge, refresh-expire, expunge",
        passive_deletes=True,
    )

    classifications = relationship(
        'Classification',
        back_populates='taxonomy',
        cascade='save-update, merge, refresh-expire, expunge',
        passive_deletes=True,
        order_by="Classification.created_at",
    )

GroupTaxonomy = join_model("group_taxonomy", Group, Taxonomy)

def get_taxonomy_usable_by_user(taxonomy_id, user_or_token):
    return (
        Taxonomy.query.filter(Taxonomy.id == taxonomy_id)
        .filter(Taxonomy.groups.any(Group.id.in_([g.id for g in user_or_token.groups])))
        .all()
    )

Taxonomy.get_taxonomy_usable_by_user = get_taxonomy_usable_by_user

class Comment(Base):

    text = sa.Column(sa.String, nullable=False)
    ctype = sa.Column(
        sa.Enum('text', 'redshift', name='comment_types', validate_strings=True)
    )

    attachment_name = sa.Column(sa.String, nullable=True)
    attachment_type = sa.Column(sa.String, nullable=True)
    attachment_bytes = sa.Column(sa.types.LargeBinary, nullable=True)

    origin = sa.Column(sa.String, nullable=True)

    author = sa.Column(sa.String, nullable=False)

    obj_id = sa.Column(
        sa.ForeignKey('objs.id', ondelete='CASCADE'), nullable=False, index=True
    )
    obj = relationship('Obj', back_populates='comments')
    groups = relationship(
        "Group",
        secondary="group_comments",
        cascade="save-update, merge, refresh-expire, expunge",
        passive_deletes=True,
    )

    def construct_author_info_dict(self):
        user = User.query.filter(User.username == self.author).first()
        return {
            field: getattr(user, field)
            for field in ('username', 'first_name', 'last_name', 'gravatar_url')
        }

    @classmethod
    def get_if_owned_by(cls, ident, user, options=[]):
        comment = cls.query.options(options).get(ident)

        if comment is not None and not comment.is_owned_by(user):
            raise AccessError('Insufficient permissions.')

        # Grab basic author info for the comment
        comment.author_info = comment.construct_author_info_dict()

        return comment

GroupComment = join_model("group_comments", Group, Comment)

class Classification(Base):
    classification = sa.Column(sa.String, nullable=False)
    taxonomy_id = sa.Column(
        sa.ForeignKey('taxonomies.id', ondelete='CASCADE'), nullable=False, index=True
    )
    taxonomy = relationship('Taxonomy', back_populates='classifications')
    probability = sa.Column(
        sa.Float,
        doc='User-assigned probability of belonging ' 'to this class',
        nullable=True,
    )

    author_id = sa.Column(
        sa.ForeignKey('users.id', ondelete='CASCADE'), nullable=False, index=True
    )
    author = relationship('User')
    author_name = sa.Column(sa.String, nullable=False)
    obj_id = sa.Column(
        sa.ForeignKey('objs.id', ondelete='CASCADE'), nullable=False, index=True
    )
    obj = relationship('Obj', back_populates='classifications')
    groups = relationship(
        "Group",
        secondary="group_classifications",
        cascade="save-update, merge, refresh-expire, expunge",
        passive_deletes=True,
    )

GroupClassifications = join_model("group_classifications", Group, Classification)

class Photometry(Base, ha.Point):
    __tablename__ = 'photometry'
    mjd = sa.Column(sa.Float, nullable=False, doc='MJD of the observation.')
    flux = sa.Column(
        sa.Float,
        doc='Flux of the observation in µJy. '
        'Corresponds to an AB Zeropoint of 23.9 in all '
        'filters.',
    )
    fluxerr = sa.Column(
        sa.Float, nullable=False, doc='Gaussian error on the flux in µJy.'
    )
    filter = sa.Column(
        allowed_bandpasses,
        nullable=False,
        doc='Filter with which the observation was taken.',
    )

    ra_unc = sa.Column(sa.Float, doc="Uncertainty of ra position [arcsec]")
    dec_unc = sa.Column(sa.Float, doc="Uncertainty of dec position [arcsec]")

    original_user_data = sa.Column(
        JSONB,
        doc='Original data passed by the user '
        'through the PhotometryHandler.POST '
        'API or the PhotometryHandler.PUT '
        'API. The schema of this JSON '
        'validates under either '
        'schema.PhotometryFlux or schema.PhotometryMag '
        '(depending on how the data was passed).',
    )
    altdata = sa.Column(JSONB)
    upload_id = sa.Column(sa.String, nullable=False, default=lambda: str(uuid.uuid4()))
    alert_id = sa.Column(sa.BigInteger, nullable=True, unique=True)

    obj_id = sa.Column(
        sa.ForeignKey('objs.id', ondelete='CASCADE'), nullable=False, index=True
    )
    obj = relationship('Obj', back_populates='photometry')
    groups = relationship(
        "Group",
        secondary="group_photometry",
        back_populates="photometry",
        cascade="save-update, merge, refresh-expire, expunge",
        passive_deletes=True,
    )
    instrument_id = sa.Column(
        sa.ForeignKey('instruments.id'), nullable=False, index=True
    )
    instrument = relationship('Instrument', back_populates='photometry')
    thumbnails = relationship('Thumbnail', passive_deletes=True)

    @hybrid_property
    def mag(self):
        if self.flux is not None and self.flux > 0:
            return -2.5 * np.log10(self.flux) + PHOT_ZP
        else:
            return None

    @hybrid_property
    def e_mag(self):
        if self.flux is not None and self.flux > 0 and self.fluxerr > 0:
            return (2.5 / np.log(10)) * (self.fluxerr / self.flux)
        else:
            return None

    @mag.expression
    def mag(cls):
        return sa.case(
            [
                (
                    sa.and_(cls.flux != None, cls.flux > 0),  # noqa
                    -2.5 * sa.func.log(cls.flux) + PHOT_ZP,
                )
            ],
            else_=None,
        )

    @e_mag.expression
    def e_mag(cls):
        return sa.case(
            [
                (
                    sa.and_(
                        cls.flux != None, cls.flux > 0, cls.fluxerr > 0
                    ),  # noqa: E711
                    2.5 / sa.func.ln(10) * cls.fluxerr / cls.flux,
                )
            ],
            else_=None,
        )

    @hybrid_property
    def jd(self):
        return self.mjd + 2_400_000.5

    @hybrid_property
    def iso(self):
        return arrow.get((self.mjd - 40_587.5) * 86400.0)

    @iso.expression
    def iso(cls):
        # converts MJD to unix timestamp
        local = sa.func.to_timestamp((cls.mjd - 40_587.5) * 86400.0)
        return sa.func.timezone('UTC', local)

    @hybrid_property
    def snr(self):
        return self.flux / self.fluxerr if self.flux and self.fluxerr else None

    @snr.expression
    def snr(self):
        return self.flux / self.fluxerr

GroupPhotometry = join_model("group_photometry", Group, Photometry)

class Spectrum(Base):
    __tablename__ = 'spectra'
    # TODO better numpy integration
    wavelengths = sa.Column(NumpyArray, nullable=False)
    fluxes = sa.Column(NumpyArray, nullable=False)
    errors = sa.Column(NumpyArray)

    obj_id = sa.Column(
        sa.ForeignKey('objs.id', ondelete='CASCADE'), nullable=False, index=True
    )
    obj = relationship('Obj', back_populates='spectra')
    observed_at = sa.Column(sa.DateTime, nullable=False)
    origin = sa.Column(sa.String, nullable=True)
    # TODO program?
    instrument_id = sa.Column(
        sa.ForeignKey('instruments.id', ondelete='CASCADE'), nullable=False, index=True
    )
    instrument = relationship('Instrument', back_populates='spectra')
    groups = relationship(
        "Group",
        secondary="group_spectra",
        back_populates="spectra",
        cascade="save-update, merge, refresh-expire, expunge",
        passive_deletes=True,
    )

    @classmethod
    def from_ascii(cls, filename, obj_id, instrument_id, observed_at):
        data = np.loadtxt(filename)
        if data.shape[1] != 2:  # TODO support other formats
            raise ValueError(f"Expected 2 columns, got {data.shape[1]}")

        return cls(
            wavelengths=data[:, 0],
            fluxes=data[:, 1],
            obj_id=obj_id,
            instrument_id=instrument_id,
            observed_at=observed_at,
        )

GroupSpectrum = join_model("group_spectra", Group, Spectrum)

# def format_public_url(context):
#    """TODO migrate this to broker tools"""
#    file_uri = context.current_parameters.get('file_uri')
#    if file_uri is None:
#        return None
#    elif file_uri.startswith('s3'):  # TODO is this reliable?
#        raise NotImplementedError
#    elif file_uri.startswith('http://'): # TODO is this reliable?
#        return file_uri
#    else:  # local file
#        return '/' + file_uri.lstrip('./')

class FollowupRequest(Base):
    requester = relationship(User, back_populates='followup_requests')
    requester_id = sa.Column(
        sa.ForeignKey('users.id', ondelete='CASCADE'), nullable=False, index=True
    )
    obj = relationship('Obj', back_populates='followup_requests')
    obj_id = sa.Column(
        sa.ForeignKey('objs.id', ondelete='CASCADE'), nullable=False, index=True
    )
    instrument = relationship(Instrument, back_populates='followup_requests')
    instrument_id = sa.Column(
        sa.ForeignKey('instruments.id'), nullable=False, index=True
    )
    start_date = sa.Column(ArrowType, nullable=False)
    end_date = sa.Column(ArrowType, nullable=False)
    filters = sa.Column(psql.ARRAY(sa.String), nullable=True)
    exposure_time = sa.Column(sa.String, nullable=True)
    priority = sa.Column(sa.Enum('1', '2', '3', '4', '5', name='priority'))
    editable = sa.Column(sa.Boolean, nullable=False, default=True)
    status = sa.Column(sa.String(), nullable=False, default="pending")

User.followup_requests = relationship('FollowupRequest', back_populates='requester')

class Thumbnail(Base):
    # TODO delete file after deleting row
    type = sa.Column(
        thumbnail_types, doc='Thumbnail type (e.g., ref, new, sub, dr8, ...)'
    )
    file_uri = sa.Column(sa.String(), nullable=True, index=False, unique=False)
    public_url = sa.Column(sa.String(), nullable=True, index=False, unique=False)
    origin = sa.Column(sa.String, nullable=True)
    photometry_id = sa.Column(
        sa.ForeignKey('photometry.id', ondelete='CASCADE'), nullable=False, index=True
    )
    photometry = relationship('Photometry', back_populates='thumbnails')
    obj = relationship(
        'Obj',
        back_populates='thumbnails',
        uselist=False,
        secondary='photometry',
        passive_deletes=True,
    )

class ObservingRun(Base):

    instrument_id = sa.Column(
        sa.ForeignKey('instruments.id', ondelete='CASCADE'), nullable=False, index=True
    )
    instrument = relationship(
        'Instrument',
        cascade='save-update, merge, refresh-expire, expunge',
        uselist=False,
        back_populates='observing_runs',
    )

    # name of the PI
    pi = sa.Column(sa.String)
    observers = sa.Column(sa.String)

    sources = relationship(
        'Obj',
        secondary='join(ClassicalAssignment, Obj)',
        cascade='save-update, merge, refresh-expire, expunge',
        passive_deletes=True,
    )

    # let this be nullable to accommodate external groups' runs
    group = relationship('Group', back_populates='observing_runs')
    group_id = sa.Column(
        sa.ForeignKey('groups.id', ondelete='CASCADE'), nullable=True, index=True
    )

    # the person who uploaded the run
    owner = relationship('User', back_populates='observing_runs')
    owner_id = sa.Column(
        sa.ForeignKey('users.id', ondelete='CASCADE'), nullable=False, index=True
    )

    assignments = relationship('ClassicalAssignment', passive_deletes=True)
    calendar_date = sa.Column(sa.Date, nullable=False, index=True)

    @property
    def _calendar_noon(self):
        observer = self.instrument.telescope.observer
        year = self.calendar_date.year
        month = self.calendar_date.month
        day = self.calendar_date.day
        hour = 12
        noon = datetime(
            year=year, month=month, day=day, hour=hour, tzinfo=observer.timezone
        )
        noon = noon.astimezone(timezone.utc).timestamp()
        noon = ap_time.Time(noon, format='unix')
        return noon

    @property
    def sunset(self):
        return self.instrument.telescope.observer.sun_set_time(
            self._calendar_noon, which='next'
        )

    @property
    def sunrise(self):
        return self.instrument.telescope.observer.sun_rise_time(
            self._calendar_noon, which='next'
        )

    @property
    def twilight_evening_nautical(self):
        return self.instrument.telescope.observer.twilight_evening_nautical(
            self._calendar_noon, which='next'
        )

    @property
    def twilight_morning_nautical(self):
        return self.instrument.telescope.observer.twilight_morning_nautical(
            self._calendar_noon, which='next'
        )

    @property
    def twilight_evening_astronomical(self):
        return self.instrument.telescope.observer.twilight_evening_astronomical(
            self._calendar_noon, which='next'
        )

    @property
    def twilight_morning_astronomical(self):
        return self.instrument.telescope.observer.twilight_morning_astronomical(
            self._calendar_noon, which='next'
        )

User.observing_runs = relationship(
    'ObservingRun', cascade='save-update, merge, refresh-expire, expunge'
)

class ClassicalAssignment(Base):
    requester_id = sa.Column(
        sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
    )
    requester = relationship(
        "User", back_populates="assignments", foreign_keys=[requester_id]
    )

    last_modified_by_id = sa.Column(
        sa.ForeignKey("users.id", ondelete="SET NULL"),
        nullable=True,
        default=None,
        index=True,
    )
    last_modified_by = relationship("User", foreign_keys=[last_modified_by_id])

    obj = relationship("Obj", back_populates="assignments")
    obj_id = sa.Column(
        sa.ForeignKey('objs.id', ondelete='CASCADE'), nullable=False, index=True
    )

    comment = sa.Column(sa.String())
    status = sa.Column(sa.String(), nullable=False, default="pending")
    priority = sa.Column(followup_priorities, nullable=False)

    run = relationship('ObservingRun', back_populates='assignments')
    run_id = sa.Column(
        sa.ForeignKey('observingruns.id', ondelete='CASCADE'),
        nullable=False,
        index=True,
    )

    @hybrid_property
    def instrument(self):
        return self.run.instrument

    @property
    def rise_time(self):
        """The time at which the object rises on this run."""
        observer = self.instrument.telescope.observer
        target = self.obj.target
        return observer.target_rise_time(
            self.run.sunset, target, which='next', horizon=30 * u.degree
        )

    @property
    def set_time(self):
        """The time at which the object sets on this run."""
        observer = self.instrument.telescope.observer
        target = self.obj.target
        return observer.target_set_time(
            self.rise_time, target, which='next', horizon=30 * u.degree
        )

User.assignments = relationship(
    "ClassicalAssignment",
    back_populates="requester",
    foreign_keys="ClassicalAssignment.requester_id",
)

schema.setup_schema()

max_stars_repo_name: dannygoldstein/skyportal

max_issues_repo_name: dannygoldstein/skyportal

max_forks_repo_name: dannygoldstein/skyportal

reason: import numpy,from astropy


text:
import pandas as pd
import numpy as np
import scipy.sparse as sparse
from xgboost import DMatrix
from copy import deepcopy
from sklearn.preprocessing import StandardScaler, MinMaxScaler

def drop_nan_bulk(data, how='all', axes=[0, 1]):
    for key in data:
        for axis in axes:
            # axis=0: rows, axis=1: cols
            # how=all: drop if all values are nan, how=any: drop if any value is nan
            data[key].dropna(axis=axis, how=how, inplace=True)
    return data

def fill_nan_bulk(data, fill='const'):
    if fill in ['mean', 'median', 'mode', 'min', 'max']:
        for key in data:
            data[key].fillna(-1, inplace=True)
            for col in data[key].select_dtypes(include=[np.number]):
                if fill == 'mean':
                    data[key].loc[:, col].replace(-1, data[key].loc[:, col].mean(), inplace=True)
                elif fill == 'median':
                    data[key].loc[:, col].replace(-1, data[key].loc[:, col].median(), inplace=True)
                elif fill == 'mode':
                    data[key].loc[:, col].replace(-1, data[key].loc[:, col].mode(), inplace=True)
                elif fill == 'min':
                    data[key].loc[:, col].replace(-1, data[key].loc[:, col].min(), inplace=True)
                elif fill == 'max':
                    data[key].loc[:, col].replace(-1, data[key].loc[:, col].max(), inplace=True)
    elif fill == 'const':
        for key in data:
            data[key].fillna(-1, inplace=True)
    elif type(fill) == [int, str]:
        for key in data:
            data[key].fillna(fill, inplace=True)
    return data

def scale_minmax_bulk(scale, data):
    for key in data:
        data[key] = pd.DataFrame(MinMaxScaler(scale).fit_transform(data[key]), index=data[key].index, columns=data[key].columns)
    return data

def scale_std_bulk(data, mean=True, std=True):
    for key in data:
        data[key] = pd.DataFrame(StandardScaler(with_mean=mean, with_std=std).fit_transform(data[key]), index=data[key].index, columns=data[key].columns)
    return data

def remove_duplicates_bulk(data):
    for key in data:
        print(key, 'rows')
        data[key] = data[key].drop_duplicates() # rows
        print(key, 'cols')
        data[key] = data[key].T.drop_duplicates().T # columns - very slow, try ndarray.swapaxes() ?
    return data

def convert_numeric_bulk(data, dtype=None, strict=False):
    for k in data:
        if dtype != None:
            if strict:
                data[k] = data[k].astype(dtype)
            else:
                for col in data[key].select_dtypes(exclude=[np.number]):
                    data[k].loc[:, col] = data[k].loc[:, col].astype(dtype)
                # for col in data[k]:
                #     if data[k][col].dtype not in ['bool', 'int64', 'float64']:
                #         data[k].loc[:, col] = data[k].loc[:, col].astype(dtype)
        else:
            data[k] = data[k].convert_objects(convert_dates=False, convert_numeric=True, convert_timedeltas=False)
    return data

def to_dmatrix(data, labels=None):
    if type(data) in [pd.DataFrame, pd.Series]:
        if labels != None and type(labels) in [pd.DataFrame, pd.Series]:
            return DMatrix(data.values, labels.values)
        else:
            return DMatrix(data.values)
    else:
        if labels != None:
            return DMatrix(data, labels)
        else:
            return DMatrix(data)

def to_sparse(data, kind='csr', fill_value=0):
    # add csc, coo variants?
    if type(data) in [pd.DataFrame, pd.Series]:
        data = data.to_sparse(fill_value)
    elif type(data) == np.ndarray:
        data = sparse.csr_matrix(data)
    return data

def to_dense(data):
    # add csc, coo variants?
    if type(data) in [pd.DataFrame, pd.Series]:
        data = data.to_dense()
    else:
        data = pd.DataFrame(data.todense())
    return data

def join(data, joins):
    for key1, key2, col in joins:
        if type(col) == str:
            if col in data[key1] or col in data[key2]:
                proceed = True
            else:
                proceed = False
        else:
            for e in col:
                if e in data[key1] or e in data[key2]:
                    proceed = True
                else:
                    proceed = False
                    break
        if proceed:
            data[key2].columns = [key2+'_'+c for c in data[key2].columns]
            try:
                if type(col) == str:
                    right_on = key2 + '_' + col
                else:
                    right_on = [key2 + '_' + c for c in col]
                data[key1] = pd.merge(data[key1], data[key2], how='left', left_on=col, right_on=right_on, sort=False)
            except KeyError:
                data[key1] = data[key1].join(data[key2], on=col)

    for key1, key2, col in joins:
        # if col in data[key1]:
        #     data[key1].drop(col, axis=1, inplace=True)
        if key2 in data:
            del data[key2]
    return data

ext: py

max_stars_repo_name: jorahn/icy

max_issues_repo_name: jorahn/icy

max_forks_repo_name: rcs-analytics/icy

reason: import numpy,import scipy


text:
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import sys
from playground.Non_DAG.utils.tools import debugPrinter
from core.scheduler import SchedulerOperation

class Node(object):
    def __init__(self, observation, action, reward, clock):
        self.observation = observation
        self.action = action
        self.reward = reward
        self.clock = clock

class RLAlgorithm(object):
    def __init__(self, agent, feature_size, features_extract_normalize_func, update_step_num, iter_num):
        super().__init__()
        self.agent = agent
        self.features_extract_normalize_func = features_extract_normalize_func
        self.feature_size = feature_size
        self.device = self.agent.device
        self.iter_num = iter_num

        self.entropy = 0
        self.current_step = 0
        self.update_step_num = update_step_num
        self.log_probs = []
        self.values    = []
        self.rewards   = []
        self.masks     = []

        # self.current_trajectory = [] # (state, action, reward, time)

    def extract_features(self, valid_pairs):
        features = []
        for machine, task, _ in valid_pairs:
            # 原来: 6维向量[machine.cpu, machine.memory, task.cpu, task.memory, task.duration, task.waiting_task_instances_number]
            # 修改: 12维向量[machine.cpu, machine.memory, machine.disk, machine.gpu, machine.gpu_memory, machine.gpu_type,
            #               task.task_config.cpu, task.task_config.memory,
            #               task.task_config.gpu, task.task_config.gpu_memory, task.task_config.gpu_type_require,
            #               task.waiting_task_instances_number]
            current_feature = machine.normal_feature + self.features_extract_normalize_func(task, machine)
            assert len(current_feature) == self.feature_size
            features.append(current_feature)
        debugPrinter(__file__, sys._getframe(), "检查feature shape: {}".format(len(features)))
        return features

    def check_rl_var(self):
        debugPrinter(__file__, sys._getframe(), "current_step: {}; entropy: {}; log_probs: {}; values: {}; rewards: {}; masks: {}"
                    .format(self.current_step, self.entropy, len(self.log_probs), len(self.values), len(self.rewards), len(self.masks)))

    def reset_rl(self):
        self.entropy = 0
        self.log_probs = []
        self.values    = []
        self.rewards   = []
        self.masks     = []
        self.current_step = 0

    def compute_returns(self, next_value, rewards, masks, gamma): # 没问题
        R = next_value
        returns = []
        for step in reversed(range(len(rewards))):
            R = rewards[step] + gamma * R * masks[step]
            returns.insert(0, R)
        return returns

    def update_brain(self, all_candidates, clock):
        if self.current_step == self.update_step_num and len(all_candidates) > 0: # TODO: 只有在笛卡尔积结果不为空时，才会更新brain
            next_features = self.extract_features(all_candidates)
            next_features = torch.tensor(next_features, dtype=torch.float32).to(self.device)

            next_value = self.agent.brain.critic_network(next_features)
            final_returns = self.compute_returns(next_value, self.rewards, self.masks, self.agent.gamma)

            self.log_probs = torch.cat(self.log_probs)
            final_returns  = torch.cat(final_returns).detach()
            self.values    = torch.cat(self.values)

            advantage = final_returns - self.values

            actor_loss  = -(self.log_probs * advantage.detach()).mean() - 0.001 * self.entropy
            critic_loss = advantage.pow(2).mean()

            self.agent.log('actor_loss', actor_loss, clock * self.iter_num)
            self.agent.log('critic_loss', critic_loss, clock * self.iter_num)

            self.agent.actor_network_optimizer.zero_grad()
            actor_loss.backward(retain_graph=True) # 保留backward后的中间参数。
            self.agent.actor_network_optimizer.step()

            self.agent.value_network_optimizer.zero_grad()
            critic_loss.backward()
            self.agent.value_network_optimizer.step()

            self.reset_rl()

    def __call__(self, cluster, clock, is_last_step=False):
        machines = cluster.machines
        tasks = cluster.tasks_which_has_waiting_instance
        all_candidates = []
        candidate_task = None
        candidate_machine = None
        schedule_candidate_task_instance = None

        for machine in machines:
            for task in tasks:
                if machine.accommodate(task):
                    all_candidates.append((machine, task, task.waiting_task_instances[0]))

        self.update_brain(all_candidates, clock)

        # TODO: 这里应该过滤掉没有任务存在时候的情况，此时不应该判定为集群资源不足，而是一种正常的逻辑?
        # TODO: 上面的这种情况下，按照我们对集群优化的理解，应该是需要执行减少集群资源节点的操作，所以这里是需要商榷的

        if len(all_candidates) == 0:
            reward = torch.FloatTensor([-1]).to(self.device) # TODO: 待设计
            sub1done = torch.FloatTensor([1 - is_last_step]).to(self.device)
            value = torch.FloatTensor([0]).to(self.device)
            log_prob = torch.FloatTensor([0]).to(self.device)
            entropy = 0

            operator_index = SchedulerOperation.OVER_SCHEDULER
        else:
            # 这里问题是: 为什么每次都只是传入state的局部观察observation，而不是整个state? 用于压缩空间?
            features = self.extract_features(all_candidates)
            features = torch.tensor(features, dtype=torch.float32).to(self.device)
            debugPrinter(__file__, sys._getframe(), "所有候选集构成的维度: {0}".format(features.shape))

            dist = self.agent.brain.actor_network(features)
            action = dist.sample() # 采样Action，单值Tensor
            log_prob = dist.log_prob(action) # 策略目标函数 J(theta) ln Π(a|s; theta) = log p(a|Π^{theta}(s))
            entropy = dist.entropy().mean() # 计算熵的平均值，用于带熵正则的策略学习中
            reward = torch.FloatTensor([1]).to(self.device) # TODO: 待设计
            sub1done = torch.FloatTensor([1 - False]).to(self.device)

            value = self.agent.brain.critic_network(features)

            candidate_machine = all_candidates[action][0]
            candidate_task = all_candidates[action][1]
            schedule_candidate_task_instance = all_candidates[action][2]

            operator_index = SchedulerOperation.TASK_INSTANCE_SCHEDULER_IN

            debugPrinter(__file__, sys._getframe(), "Action: (machine: {0}; task: {1})".format(candidate_machine.id, candidate_task.task_index))

        self.entropy += entropy
        self.log_probs.append(log_prob)
        self.values.append(value)
        self.rewards.append(reward)
        self.masks.append(sub1done)
        self.current_step += 1

        return operator_index, candidate_machine, candidate_task, schedule_candidate_task_instance

reason: import numpy


text:
#!/usr/bin/env python
##!/d1/local/bin/python
##!/d1/local/Canopy/bin/python
##!/d1/local/enthought/bin/python
##!/usr/bin/env python
##!/d1/local/src/Python-2.6.2/python
# used to be /usr/bin/env python

""""""

# ============================================================================== #
#                                                                                #
#   (c) Copyright, 2011 University Corporation for Atmospheric Research (UCAR).  #
#       All rights reserved.                                                     #
#                                                                                #
#       File: $RCSfile: view_state_shape_file.py,v $                                           #
#       Version: $Revision: 1.2 $  Dated: $Date: 2014/06/18 22:01:09 $           #
#                                                                                #
# ============================================================================== #

import httplib
import xml.dom.minidom

import log_msg
import os
import sys
from optparse import OptionParser
#import psycopg2 as pg
from osgeo import ogr
from osgeo import osr
from netCDF4 import Dataset
import numpy
import pyproj
import math
import datetime
import re

#cdl = "/d2/vii/data/static/cdl/road_segments.cdl"
GRC_EARTH_RADIUS = 6371
#NV_INTERSTATES = "Interstate 80"
#MN_INTERSTATES = ["Interstate 35", "Interstate 35e", "Interstate 35w", "Interstate 535", "Interstate 90", "Interstate 94", "Interstate 494",  "Interstate 694"]

def dist(lat1,lon1,lat2,lon2):

    """Great circle distance"""

    rlat1 = math.radians(lat1)
    rlat2 = math.radians(lat2)
    rlon1 = math.radians(lon1)
    rlon2 = math.radians(lon2)
    delta_lat = rlat2 - rlat1;
    delta_lon = rlon2 - rlon1;

    sin_delta_lat2 = math.sin(delta_lat / 2.0)
    sin_delta_lat2_sq = sin_delta_lat2 * sin_delta_lat2
    sin_delta_lon2 = math.sin(delta_lon / 2.0)
    sin_delta_lon2_sq = sin_delta_lon2 * sin_delta_lon2

    temp = sin_delta_lat2_sq + math.cos(rlat1) * math.cos(rlat2) * sin_delta_lon2_sq
    distance = GRC_EARTH_RADIUS * 2 * math.atan2(math.sqrt(temp), math.sqrt(1 - temp))

    return distance

def seg_length(seg, distances):
    t_dist = 0.0
    #print "len(seg)-1: ", len(seg)-1
    for a in range(len(seg)-1):
        pt1 = list(seg[a])
        pt2 = list(seg[a+1])
        if pt1[0] > 0:
            pt1[0] *= -1
        if pt2[0] > 0:
            pt2[0] *= -1
        d = dist(pt1[1],pt1[0],pt2[1],pt2[0])

        #print "a: ", a
        #print "distance: ", d
        distances.append(d)
        #print "distances: ", distances
        t_dist += d
        #print "t_dist: ", t_dist
    return t_dist

def set_mid_point(seg, total_distance, distances, mid_point_lats, mid_point_lons):
    half_total_distance = total_distance / 2.0
    #print "half_total_distance: ", half_total_distance
    partial_distance = 0
    previous_partial_distance = 0;
    mid_index = 0
    #print "distances: ", distances
    #print "seg: ", seg
    #print
    for b in range(len(seg)-1):
        #print "b: ", b
        previous_partial_distance = partial_distance
        partial_distance += distances[b]
        #print "partial_distance, half_total_distance: %f, %f" % (partial_distance, half_total_distance)
        #print "previous_partial_distance: ", previous_partial_distance
        if (partial_distance >= half_total_distance):
            #print "partial_distance: ", partial_distance
            if (half_total_distance - previous_partial_distance) < (partial_distance - half_total_distance):
                mid_index = b
            else:
                mid_index = b+1
            break;
    mid_point_lats.append(seg[mid_index][1])
    mid_point_lons.append(seg[mid_index][0])
    #print "mid_point_lat: ", seg[mid_index][1]
    #print "mid_point_lon: ", seg[mid_index][0]

def segmentize(line, seg_len, mid_point_lats, mid_point_lons):
    """Pull out the points in the polyline. Note that in the GIS world a polyline is defined as a set of lines that may not be connected"""
    #print "line: ", line
    #print
    segs = []
    #distances = []
    #total_distance = 0
    n = line.GetPointCount()
    #print "line.GetPointCount(): ", n
    seg = []
    # Iterate over the number of points in the polyline
    for i in range(n):
        distances = []
        total_distance = 0
        pt1 = line.GetPoint_2D(i)
        #print "i: ", i
        #print "pt1: ", pt1
        #print
        seg.append(pt1)
        d = seg_length(seg, distances)
        #print "seg: ", seg
        #print
        #print "distance: ", d
        # Check to see if the total length of the points so far is greater than the specified segment length
        if d >= seg_len:                     # 1.6 km for 1 mile segments
            #print "Total distance of segment in kilometers, in miles: %f, %f" % (d, d*0.6214)
            #print "seg d>= seg_len: ", seg
            #print
            # If the desired segment length (or greater) has been reached, append the point to the list of segments
            total_distance = d
            #print "total_distance: ", total_distance
            segs.append(seg)
            set_mid_point(seg, total_distance, distances, mid_point_lats, mid_point_lons)
            seg = [pt1]
    #print "out of for loop"
    return segs

def set_aux_id(seg_id, state):
    if (len(str(seg_id)) == 1):
        if state == "NV":
            aux_id = "N0000%d" % seg_id
        elif state == "MN":
            aux_id = "M0000%d" % seg_id
        elif state == "MI":
            aux_id = "I0000%d" % seg_id
        elif state == "CO":
            aux_id = "C0000%d" % seg_id
        return aux_id;
    elif (len(str(seg_id)) == 2):
        if state == "NV":
            aux_id = "N000%d" % seg_id
        elif state == "MN":
            aux_id = "M000%d" % seg_id
        elif state == "MI":
            aux_id = "I000%d" % seg_id
        elif state == "CO":
            aux_id = "C000%d" % seg_id
        return aux_id;
    elif (len(str(seg_id)) == 3):
        if state == "NV":
            aux_id = "N00%d" % seg_id
        elif state == "MN":
            aux_id = "M00%d" % seg_id
        elif state == "MI":
            aux_id = "I00%d" % seg_id
        elif state == "CO":
            aux_id = "C00%d" % seg_id
        return aux_id;
    elif (len(str(seg_id)) == 4):
        if state == "NV":
            aux_id = "N0%d" % seg_id
        elif state == "MN":
            aux_id = "M0%d" % seg_id
        elif state == "MI":
            aux_id = "I0%d" % seg_id
        elif state == "CO":
            aux_id = "C0%d" % seg_id
        return aux_id;
    elif (len(str(seg_id)) == 5):
        if state == "NV":
            aux_id = "N%d" % seg_id
        elif state == "MN":
            aux_id = "M%d" % seg_id
        elif state == "MI":
            aux_id = "I%d" % seg_id
        elif state == "CO":
            aux_id = "C%d" % seg_id
        return aux_id;
    elif (len(str(seg_id)) >= 6):
        logg.write_error("Too many segment ids. Unable to set aux_id.")
        exit(1);

def main():

    usage_str = "%prog shp_file prj_file seg_len state_abbrev"
    parser = OptionParser(usage = usage_str)
    parser.add_option("-l", "--log", dest="log", help="write log messages to specified file")

    (options, args) = parser.parse_args()

    if len(args) < 3:
        parser.print_help()
        sys.exit(2)

    shp_file = args[0]
    prj_file = args[1]
    seg_len = args[2]
    state = args[3]

    try:
        # seg_len should be in kilometers
        seg_len = float(seg_len)
    except:
        print "couldn't parse %s as segment length" % seg_len
        sys.exit(2)

    if options.log:
        logg = log_msg.LogMessage(options.log)
    else:
        logg = log_msg.LogMessage("")

    logg.write_starting()

    pt_names = []
    lats = []
    lons = []
    mid_point_lats = []
    mid_point_lons = []
    mid_point_latitudes = []
    mid_point_longitudes = []
    elevs = []
    ids = []
    aux_ids = []
    all_names = {}

    # Initialize ogr shape library
    ogr.RegisterAll()

    # Open shape file
    #ds = ogr.Open("/home/vii/mn_road_segments_from_jakin/Current_ALL_plowroutes.shp")
    logg.write_info("open %s" % shp_file)
    ds = ogr.Open(shp_file)

    # Read in the prj file
    #wkt = open("/home/vii/mn_road_segments_from_jakin/Current_ALL_plowroutes.prj").read()
    wkt = open(prj_file).read()

    # Shape file has a concept of layers. There can be multiple layers
    # such as roads, railroads, etc.  Most of the time its one layer
    # and here are the roads Get 1st layer
    root = ds.GetLayerByIndex(0)

    srs = osr.SpatialReference()
    srs.ImportFromESRI([wkt])

    proj = osr.SpatialReference()
    proj.ImportFromProj4(srs.ExportToProj4())

    # default projection
    def_proj = osr.SpatialReference()
    def_proj.ImportFromEPSG(4326)
    #seg_id = 1
    seg_id = 0

    logg.write_info("using %.5fkm as minimum segment length" % seg_len)
    logg.write_info("has %d features" % len(root))

    # Loop through all the lines in the layer
    lines = []
    road_names = []

    for i in range(len(root)):
        # Get the ith feature (road segment)
        f = root.GetFeature(i)
        #print "ith feature: ", dir(f)

        # Just get the geometry of the feature (there can be
        # additional information for the feature such as name, etc.)
        geom = f.GetGeometryRef()

        # Sometimes features have no geometry
        if geom == None:
            continue

        # Assign projection
        geom.AssignSpatialReference(proj)

        # Transform to default projection
        geom.TransformTo(def_proj)

        # Get the name of the route (junction I35 with Chicago say)
        road_name = "unknown"
        for k in range(f.GetFieldCount()):
            field_name = f.GetFieldDefnRef(k).GetName()
            #print "field_name: ", field_name
            field_string = f.GetFieldAsString(k)
            #print "field_string: ", field_string

            # Only pull out Interstate 80 at this time
            if "lname" in field_name.lower():
                lname_string = f.GetFieldAsString(k)
                #print lname_string
                if state == "NV":
                    if re.search(r"^Interstate", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                elif state == "MN":
                    if re.search(r"^Interstate", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                elif state == "MI":
                    if re.search(r"^Interstate", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                        #print lname_string
                    if re.search(r"^Chevrolet Buick Frwy", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                        #print lname_string
                    if re.search(r"^Jeffries Frwy", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                        #print lname_string
                    if re.search(r"^Edsel Ford Frwy", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                        #print lname_string
                    if re.search(r"^Chrysler Frwy", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                        #print lname_string
                    if re.search(r"^Fisher Frwy", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                        #print lname_string
                elif state == "CO":
                    if re.search(r"^US Hwy", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                        #print lname_string
                    if re.search(r"^Interstate", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                        #print lname_string
                    if re.search(r"^State Hwy", lname_string):
                        road_name = lname_string
                        road_names.append(road_name)
                        #print lname_string
                break

        # Loop over the points in the geometry
        # A point is still a geometry but here a polyline consists of a number of line segments
        c = geom.GetGeometryCount()
        #print "c: ", c
        if c > 0:
            print "c: ", c
            for j in range(c):
                line = geom.GetGeometryRef(j)

                # Only pull out lat/lon for Interstate 80 at this time
                if (road_name != "unknown"):
                    #print "road_name: ", road_name
                    #print "geom: ", line
                    lines.append(line.Clone()) # use clone because GeometryRef will be garbage collected after this loop, causing a segfault
        else:
            # Only pull out lat/lon for Interstate 80 at this time
            if (road_name != "unknown"):
                #print "road_name: ", road_name
                #print "geom: ", geom.Clone()
                lines.append(geom.Clone())
    for index, line in enumerate(lines):
        segs = segmentize(line, seg_len, mid_point_lats, mid_point_lons)
        #print "line: ", line
        #print "len(lines): ", len(lines)
        #print "len(road_names): ", len(road_names)
        tmp_aux_id = ""
        #for index, seg in enumerate(segs):
        for ind, seg in enumerate(segs):
            #print "len(seg): ", len(seg)
            #print "seg (second print): ", seg
            #print
            #print "INDEX: ", index
            #print "top level, road_names[index]: ", road_names[index]
            if not all_names.has_key(road_names[index]):
                name = road_names[index]
                all_names[road_names[index]] = 1
                name = "%s %d" % (road_names[index], all_names[road_names[index]])
                seg_id += 1
                #print "name: ", name
                #print "road_names[index]: ", road_names[index]
                #print "seg_id: ", seg_id
            else:
                name = "%s %d" % (road_names[index], all_names[road_names[index]]+1)
                all_names[road_names[index]] += 1
                seg_id += 1
                #print "has key, name: ", name
                #print "has key, road_names[index]: ", road_names[index]
                #print "has key, seg_id: ", seg_id
            name = list(name.ljust(32,'\0'))
            tmp_aux_id = set_aux_id(seg_id, state)
            aux_id = list(tmp_aux_id.ljust(7,'\0'))

            for (lon,lat) in seg:
                lats.append(lat)
                #print "lat: ", lat
                lons.append(lon)
                #print "lon: ", lon
                pt_names.append(name)
                #print "name: ", name
                ids.append(seg_id)
                #print "seg_id: ", seg_id
                mid_point_latitudes.append(mid_point_lats[ind])
                mid_point_longitudes.append(mid_point_lons[ind])
                aux_ids.append(aux_id)

    # Now have a list of different information for the line in
    # lats, lons, pt_names, ids

    logg.write_ending()

if __name__ == "__main__":
   main()

max_stars_repo_name: OSADP/Pikalert-Vehicle-Data-Translator-

max_issues_repo_name: OSADP/Pikalert-Vehicle-Data-Translator-

max_forks_repo_name: OSADP/Pikalert-Vehicle-Data-Translator-


text:
import numpy as np
import pandas as pd
from sklearn.linear_model import ElasticNetCV
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.tree import DecisionTreeRegressor
from tpot.builtins import StackingEstimator
from tpot.export_utils import set_param_recursive

# NOTE: Make sure that the outcome column is labeled 'target' in the data file
tpot_data = pd.read_csv("PATH/TO/DATA/FILE",
                        sep="COLUMN_SEPARATOR", dtype=np.float64)
features = tpot_data.drop("target", axis=1)
training_features, testing_features, training_target, testing_target = train_test_split(
    features, tpot_data["target"], random_state=42
)

# Average CV score on the training set was: -0.23721048542910758
exported_pipeline = make_pipeline(
    StackingEstimator(
        estimator=DecisionTreeRegressor(
            max_depth=1, min_samples_leaf=5, min_samples_split=15
        )
    ),
    StackingEstimator(estimator=ElasticNetCV(
        l1_ratio=0.9500000000000001, tol=0.1)),
    ElasticNetCV(l1_ratio=0.45, tol=0.001),
)
# Fix random state for all the steps in exported pipeline
set_param_recursive(exported_pipeline.steps, "random_state", 42)

exported_pipeline.fit(training_features, training_target)
results = exported_pipeline.predict(testing_features)


text:
# Copyright 2020 Huawei Technologies Co., Ltd
#
# 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.
# ==============================================================================
"""
graphdata.py supports loading graph dataset for GNN network training,
and provides operations related to graph data.
"""
import numpy as np
from mindspore._c_dataengine import Graph
from mindspore._c_dataengine import Tensor

from .validators import check_gnn_get_all_nodes, check_gnn_get_all_neighbors, check_gnn_get_node_feature

class GraphData:
    """
    Reads th graph dataset used for GNN training from the shared file and database.

    Args:
        dataset_file (str): One of file names in dataset.
        num_parallel_workers (int, optional): Number of workers to process the Dataset in parallel
            (default=None).
    """

    def __init__(self, dataset_file, num_parallel_workers=None):
        self._dataset_file = dataset_file
        if num_parallel_workers is None:
            num_parallel_workers = 1
        self._graph = Graph(dataset_file, num_parallel_workers)

    @check_gnn_get_all_nodes
    def get_all_nodes(self, node_type):
        """
        Get all nodes in the graph.

        Args:
            node_type (int): Specify the tpye of node.

        Returns:
            numpy.ndarray: array of nodes.

        Examples:
            >>> import mindspore.dataset as ds
            >>> data_graph = ds.GraphData('dataset_file', 2)
            >>> nodes = data_graph.get_all_nodes(0)

        Raises:
            TypeError: If `node_type` is not integer.
        """
        return self._graph.get_nodes(node_type, -1).as_array()

    @check_gnn_get_all_neighbors
    def get_all_neighbors(self, node_list, neighbor_type):
        """
        Get `neighbor_type` neighbors of the nodes in `node_list`.

        Args:
            node_list (list or numpy.ndarray): The given list of nodes.
            neighbor_type (int): Specify the tpye of neighbor.

        Returns:
            numpy.ndarray: array of nodes.

        Examples:
            >>> import mindspore.dataset as ds
            >>> data_graph = ds.GraphData('dataset_file', 2)
            >>> nodes = data_graph.get_all_nodes(0)
            >>> neighbors = data_graph.get_all_neighbors(nodes[0], 0)

        Raises:
            TypeError: If `node_list` is not list or ndarray.
            TypeError: If `neighbor_type` is not integer.
        """
        return self._graph.get_all_neighbors(node_list, neighbor_type).as_array()

    @check_gnn_get_node_feature
    def get_node_feature(self, node_list, feature_types):
        """
        Get `feature_types` feature of the nodes in `node_list`.

        Args:
            node_list (list or numpy.ndarray): The given list of nodes.
            feature_types (list or ndarray): The given list of feature types.

        Returns:
            numpy.ndarray: array of features.

        Examples:
            >>> import mindspore.dataset as ds
            >>> data_graph = ds.GraphData('dataset_file', 2)
            >>> nodes = data_graph.get_all_nodes(0)
            >>> features = data_graph.get_node_feature(nodes[0], [1])

        Raises:
            TypeError: If `node_list` is not list or ndarray.
            TypeError: If `feature_types` is not list or ndarray.
        """
        if isinstance(node_list, list):
            node_list = np.array(node_list, dtype=np.int32)
        return [t.as_array() for t in self._graph.get_node_feature(Tensor(node_list), feature_types)]

max_stars_repo_name: ythlml/mindspore

max_issues_repo_name: ythlml/mindspore

max_forks_repo_name: ythlml/mindspore


text:
"""This file contains code for use with "Think Stats",
by Allen B. Downey, available from greenteapress.com

Copyright 2014 Allen B. Downey
License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html
"""

from __future__ import print_function

import numpy as np
import sys

import nsfg
import thinkstats2

def validatePregnum(resp):
    preg = nsfg.ReadFemPreg(dct_file='data/2002FemPreg.dct', dat_file='data/2002FemPreg.dat.gz')
    preg_map = nsfg.MakePregMap(preg)
    for index, pregnum in resp.pregnum.items():
        caseid = resp.caseid[index]
        indices = preg_map[caseid]

        # check that pregnum from the respondent file equals
        # the number of records in the pregnancy file
        if len(indices) != pregnum:
            print(caseid, len(indices), pregnum)
            return False

    return True

def main(script):
    """Tests the functions in this module.

    script: string script name
    """
    resp = nsfg.ReadFemResp(dct_file='data/2002FemResp.dct', dat_file='data/2002FemResp.dat.gz')
    assert(len(resp) == 7643)
    assert(resp.pregnum.value_counts()[1] == 1267)
    assert(validatePregnum(resp))

    print('%s: All tests passed.' % script)

if __name__ == '__main__':
    main(*sys.argv)


text:
# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved.                     #
# This file is part of the AiiDA code.                                    #
#                                                                         #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file        #
# For further information please visit http://www.aiida.net               #
###########################################################################
# pylint: disable=invalid-name
"""Various utils that should be used during migrations and migrations tests because the AiiDA ORM cannot be used."""

import datetime
import errno
import os
import re

import numpy

from aiida.common import json

ISOFORMAT_DATETIME_REGEX = re.compile(r'^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d+(\+\d{2}:\d{2})?$')

def ensure_repository_folder_created(uuid):
    """Make sure that the repository sub folder for the node with the given UUID exists or create it.

    :param uuid: UUID of the node
    """
    dirpath = get_node_repository_sub_folder(uuid)

    try:
        os.makedirs(dirpath)
    except OSError as exception:
        if exception.errno != errno.EEXIST:
            raise

def put_object_from_string(uuid, name, content):
    """Write a file with the given content in the repository sub folder of the given node.

    :param uuid: UUID of the node
    :param name: name to use for the file
    :param content: the content to write to the file
    """
    ensure_repository_folder_created(uuid)
    filepath = os.path.join(get_node_repository_sub_folder(uuid), name)

    with open(filepath, 'w', encoding='utf-8') as handle:
        handle.write(content)

def get_object_from_repository(uuid, name):
    """Return the content of a file with the given name in the repository sub folder of the given node.

    :param uuid: UUID of the node
    :param name: name to use for the file
    """
    filepath = os.path.join(get_node_repository_sub_folder(uuid), name)

    with open(filepath) as handle:
        return handle.read()

def get_node_repository_sub_folder(uuid):
    """Return the absolute path to the sub folder `path` within the repository of the node with the given UUID.

    :param uuid: UUID of the node
    :return: absolute path to node repository folder, i.e `/some/path/repository/node/12/ab/c123134-a123/path`
    """
    from aiida.manage.configuration import get_profile

    uuid = str(uuid)

    repo_dirpath = os.path.join(get_profile().repository_path, 'repository')
    node_dirpath = os.path.join(repo_dirpath, 'node', uuid[:2], uuid[2:4], uuid[4:], 'path')

    return node_dirpath

def get_numpy_array_absolute_path(uuid, name):
    """Return the absolute path of a numpy array with the given name in the repository of the node with the given uuid.

    :param uuid: the UUID of the node
    :param name: the name of the numpy array
    :return: the absolute path of the numpy array file
    """
    return os.path.join(get_node_repository_sub_folder(uuid), name + '.npy')

def store_numpy_array_in_repository(uuid, name, array):
    """Store a numpy array in the repository folder of a node.

    :param uuid: the node UUID
    :param name: the name under which to store the array
    :param array: the numpy array to store
    """
    ensure_repository_folder_created(uuid)
    filepath = get_numpy_array_absolute_path(uuid, name)

    with open(filepath, 'wb') as handle:
        numpy.save(handle, array)

def delete_numpy_array_from_repository(uuid, name):
    """Delete the numpy array with a given name from the repository corresponding to a node with a given uuid.

    :param uuid: the UUID of the node
    :param name: the name of the numpy array
    """
    filepath = get_numpy_array_absolute_path(uuid, name)

    try:
        os.remove(filepath)
    except (IOError, OSError):
        pass

def load_numpy_array_from_repository(uuid, name):
    """Load and return a numpy array from the repository folder of a node.

    :param uuid: the node UUID
    :param name: the name under which to store the array
    :return: the numpy array
    """
    filepath = get_numpy_array_absolute_path(uuid, name)
    return numpy.load(filepath)

def recursive_datetime_to_isoformat(value):
    """Convert all datetime objects in the given value to string representations in ISO format.

    :param value: a mapping, sequence or single value optionally containing datetime objects
    """
    if isinstance(value, list):
        return [recursive_datetime_to_isoformat(_) for _ in value]

    if isinstance(value, dict):
        return dict((key, recursive_datetime_to_isoformat(val)) for key, val in value.items())

    if isinstance(value, datetime.datetime):
        return value.isoformat()

    return value

def dumps_json(dictionary):
    """Transforms all datetime object into isoformat and then returns the JSON."""
    return json.dumps(recursive_datetime_to_isoformat(dictionary))

max_stars_repo_name: pranavmodx/aiida-core

max_issues_repo_name: pranavmodx/aiida-core

max_forks_repo_name: lorisercole/aiida-core


text:
'''
import numpy as np

from ._lags_matrix import conv_matrix

#TODO: merge both implementations
# probably for this padded signal need to be with mean value.
__all__=['moving_average','movav']

def moving_average(vector,av_cof, mode='postw'):

    Simple Moving Average or Weighted  Moving Average

    Parameters
    -------------
    * vector: 1d ndarray,
        input vector.
    * av_cof: int or 1d ndarray,
        > if av_cof is int:averaging-window size;
        > if av_cof is 1d ndarray:weights array
                        for weighted moving average;
    * mode: string,
        mode = {'full',prew','postw','valid','same'}
        > mode = full: output with size = av_cof+N-1
        > mode = prew: output with size = N,
            padding of zeros in the first.
        > mode = postw: output with size = N,
             padding of zeros in the last.
        > mode = valid: output with size = N-av_cof+1,
             without padding of zeros.
        > mode = same: output with size = N,
             without padding of zeros at the begin and end.
     Returns
    -----------
    * smoothed array: 1d ndarray.

    av_cof = np.asarray(av_cof,dtype=int)

    if av_cof.size ==1:
        window = np.ones(av_cof)
    else:
        window = av_cof

    vector = np.asarray(vector)

    return conv_matrix(vector,mode=mode,lags=window.size) @ window/window.size
#---------------------------------------------------
def movav(vector,window_size,straight = True):

    Alternative Moving Average Implementation.

    Paramters
    -----------------
    vector: 1d ndarray,
        input vector.
    window_size: size of sliding window.
    straight: bool,
        if true, sliding in straight direction,
        if false, sliding in backward direction.

    Return
    ----------
    smoothed vector: 1d ndarray.

    vector = np.asarray(vector)
    N = vector.shape[0]
    out = np.zeros_like(vector)

    if straight:
#         if symmetry:
        for i in range(N):
            lp = min(N,i+window_size//2)
            fp = max(i-window_size - window_size//2,0)
            out[i] = np.mean(vector[fp:lp])
#         else:
#             for i in range(N):
#                 lp = min(N,i+size)
#                 out[i] = np.mean(vector[i:lp])
    else:
        for i in range(N-1,-1,-1):
            fp = N-min(N-1,i+window_size//2)-1
            lp = N-max(i-window_size - window_size//2,0)-1
            out[N-i-1] = np.mean(vector[fp:lp])
#         else:
#             for i in range(N,0,-1):
#                 lp = N-min(N,i+size)+1
#                 out[i] = np.mean(vector[N-i+1:lp])
    return out

max_stars_repo_name: diarmaidocualain/dsatools

max_issues_repo_name: Jimmy-INL/dsatools

max_forks_repo_name: Jimmy-INL/dsatools


text:
import os
from collections import deque

import dmc2gym
import gym
import numpy as np
import torch
import torch.nn.functional as F
import torchvision.transforms.functional as TF
from numpy.random import randint

def make_env(domain_name, task_name, seed=0, episode_length=1000, frame_stack=3, action_repeat=4, image_size=100,
             mode='train'):
    """Make environment for experiments"""
    assert mode in {'train', 'color_easy', 'color_hard', 'video_easy', 'video_hard'}, \
        f'specified mode "{mode}" is not supported'

    env = dmc2gym.make(
        domain_name=domain_name,
        task_name=task_name,
        seed=seed,
        visualize_reward=False,
        from_pixels=True,
        height=image_size,
        width=image_size,
        episode_length=episode_length,
        frame_skip=action_repeat
    )
    env = VideoWrapper(env, mode, seed)
    env = FrameStack(env, frame_stack)
    env = ColorWrapper(env, mode, seed)

    return env

class ColorWrapper(gym.Wrapper):
    """Wrapper for the color experiments"""

    prefix = "custom_vendor/data"

    def __init__(self, env, mode, seed=None):
        assert isinstance(env, FrameStack), 'wrapped env must be a framestack'
        gym.Wrapper.__init__(self, env)
        self._max_episode_steps = env._max_episode_steps
        self._mode = mode
        self._random_state = np.random.RandomState(seed)
        self.time_step = 0
        if 'color' in self._mode:
            self._load_colors()

    def reset(self):
        self.time_step = 0
        if 'color' in self._mode:
            self.randomize()
        elif 'video' in self._mode:
            # apply greenscreen
            setting_kwargs = {
                'skybox_rgb': [.2, .8, .2],
                'skybox_rgb2': [.2, .8, .2],
                'skybox_markrgb': [.2, .8, .2]
            }
            if self._mode == 'video_hard':
                setting_kwargs['grid_rgb1'] = [.2, .8, .2]
                setting_kwargs['grid_rgb2'] = [.2, .8, .2]
                setting_kwargs['grid_markrgb'] = [.2, .8, .2]
            self.reload_physics(setting_kwargs)
        return self.env.reset()

    def step(self, action):
        self.time_step += 1
        return self.env.step(action)

    def randomize(self):
        assert 'color' in self._mode, f'can only randomize in color mode, received {self._mode}'
        self.reload_physics(self.get_random_color())

    def _load_colors(self):
        assert self._mode in {'color_easy', 'color_hard'}
        self._colors = torch.load(f'{self.prefix}/{self._mode}.pt')

    def get_random_color(self):
        assert len(self._colors) >= 100, 'env must include at least 100 colors'
        return self._colors[self._random_state.randint(len(self._colors))]

    def reload_physics(self, setting_kwargs=None, state=None):
        from dm_control.suite import common
        domain_name = self._get_dmc_wrapper()._domain_name
        if setting_kwargs is None:
            setting_kwargs = {}
        if state is None:
            state = self._get_state()
        self._reload_physics(
            *common.settings.get_model_and_assets_from_setting_kwargs(
                domain_name + '.xml', setting_kwargs
            )
        )
        self._set_state(state)

    def get_state(self):
        return self._get_state()

    def set_state(self, state):
        self._set_state(state)

    def _get_dmc_wrapper(self):
        _env = self.env
        while not isinstance(_env, dmc2gym.wrappers.DMCWrapper) and hasattr(_env, 'env'):
            _env = _env.env
        assert isinstance(_env, dmc2gym.wrappers.DMCWrapper), 'environment is not dmc2gym-wrapped'

        return _env

    def _reload_physics(self, xml_string, assets=None):
        _env = self.env
        while not hasattr(_env, '_physics') and hasattr(_env, 'env'):
            _env = _env.env
        assert hasattr(_env, '_physics'), 'environment does not have physics attribute'
        _env.physics.reload_from_xml_string(xml_string, assets=assets)

    def _get_physics(self):
        _env = self.env
        while not hasattr(_env, '_physics') and hasattr(_env, 'env'):
            _env = _env.env
        assert hasattr(_env, '_physics'), 'environment does not have physics attribute'

        return _env._physics

    def _get_state(self):
        return self._get_physics().get_state()

    def _set_state(self, state):
        self._get_physics().set_state(state)

class FrameStack(gym.Wrapper):
    """Stack frames as observation"""

    def __init__(self, env, k):
        gym.Wrapper.__init__(self, env)
        self._k = k
        self._frames = deque([], maxlen=k)
        shp = env.observation_space.shape
        self.observation_space = gym.spaces.Box(
            low=0,
            high=1,
            shape=((shp[0] * k,) + shp[1:]),
            dtype=env.observation_space.dtype
        )
        self._max_episode_steps = env._max_episode_steps

    def reset(self):
        obs = self.env.reset()
        for _ in range(self._k):
            self._frames.append(obs)
        return self._get_obs()

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        self._frames.append(obs)
        return self._get_obs(), reward, done, info

    def _get_obs(self):
        assert len(self._frames) == self._k
        return np.concatenate(list(self._frames), axis=0)

def rgb_to_hsv(r, g, b):
    """Convert RGB color to HSV color"""
    maxc = max(r, g, b)
    minc = min(r, g, b)
    v = maxc
    if minc == maxc:
        return 0.0, 0.0, v
    s = (maxc - minc) / maxc
    rc = (maxc - r) / (maxc - minc)
    gc = (maxc - g) / (maxc - minc)
    bc = (maxc - b) / (maxc - minc)
    if r == maxc:
        h = bc - gc
    elif g == maxc:
        h = 2.0 + rc - bc
    else:
        h = 4.0 + gc - rc
    h = (h / 6.0) % 1.0
    return h, s, v

def do_green_screen(x, bg):
    """Removes green background from observation and replaces with bg; not optimized for speed"""
    assert isinstance(x, np.ndarray) and isinstance(bg, np.ndarray), 'inputs must be numpy arrays'
    assert x.dtype == np.uint8 and bg.dtype == np.uint8, 'inputs must be uint8 arrays'

    # Get image sizes
    x_h, x_w = x.shape[1:]

    # Convert to RGBA images
    im = TF.to_pil_image(torch.ByteTensor(x))
    im = im.convert('RGBA')
    pix = im.load()
    bg = TF.to_pil_image(torch.ByteTensor(bg))
    bg = bg.convert('RGBA')
    bg = bg.load()

    # Replace pixels
    for x in range(x_w):
        for y in range(x_h):
            r, g, b, a = pix[x, y]
            h_ratio, s_ratio, v_ratio = rgb_to_hsv(r / 255., g / 255., b / 255.)
            h, s, v = (h_ratio * 360, s_ratio * 255, v_ratio * 255)

            min_h, min_s, min_v = (100, 80, 70)
            max_h, max_s, max_v = (185, 255, 255)
            if min_h <= h <= max_h and min_s <= s <= max_s and min_v <= v <= max_v:
                pix[x, y] = bg[x, y]

    return np.moveaxis(np.array(im).astype(np.uint8), -1, 0)[:3]

class VideoWrapper(gym.Wrapper):
    """Green screen for video experiments"""
    prefix = "custom_vendor/data"

    def __init__(self, env, mode, seed):
        gym.Wrapper.__init__(self, env)
        self._mode = mode
        self._seed = seed
        self._random_state = np.random.RandomState(seed)
        self._index = 0
        self._video_paths = []
        if 'video' in mode:
            self._get_video_paths()
        self._num_videos = len(self._video_paths)
        self._max_episode_steps = env._max_episode_steps

    def _get_video_paths(self):
        video_dir = os.path.join(self.prefix, self._mode)
        if 'video_easy' in self._mode:
            self._video_paths = [os.path.join(video_dir, f'video{i}.mp4') for i in range(10)]
        elif 'video_hard' in self._mode:
            self._video_paths = [os.path.join(video_dir, f'video{i}.mp4') for i in range(100)]
        else:
            raise ValueError(f'received unknown mode "{self._mode}"')

    def _load_video(self, video):
        """Load video from provided filepath and return as numpy array"""
        import cv2
        cap = cv2.VideoCapture(video)
        assert cap.get(cv2.CAP_PROP_FRAME_WIDTH) >= 100, 'width must be at least 100 pixels'
        assert cap.get(cv2.CAP_PROP_FRAME_HEIGHT) >= 100, 'height must be at least 100 pixels'
        n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        buf = np.empty((n, int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), 3),
                       np.dtype('uint8'))
        i, ret = 0, True
        while (i < n and ret):
            ret, frame = cap.read()
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            buf[i] = frame
            i += 1
        cap.release()
        return np.moveaxis(buf, -1, 1)

    def _reset_video(self):
        self._index = (self._index + 1) % self._num_videos
        self._data = self._load_video(self._video_paths[self._index])

    def reset(self):
        if 'video' in self._mode:
            self._reset_video()
        self._current_frame = 0
        return self._greenscreen(self.env.reset())

    def step(self, action):
        self._current_frame += 1
        obs, reward, done, info = self.env.step(action)
        return self._greenscreen(obs), reward, done, info

    def _interpolate_bg(self, bg, size: tuple):
        """Interpolate background to size of observation"""
        bg = torch.from_numpy(bg).float().unsqueeze(0) / 255.
        bg = F.interpolate(bg, size=size, mode='bilinear', align_corners=False)
        return (bg * 255.).byte().squeeze(0).numpy()

    def _greenscreen(self, obs):
        """Applies greenscreen if video is selected, otherwise does nothing"""
        if 'video' in self._mode:
            bg = self._data[self._current_frame % len(self._data)]  # select frame
            bg = self._interpolate_bg(bg, obs.shape[1:])  # scale bg to observation size
            return do_green_screen(obs, bg)  # apply greenscreen
        return obs

    def apply_to(self, obs):
        """Applies greenscreen mode of object to observation"""
        obs = obs.copy()
        channels_last = obs.shape[-1] == 3
        if channels_last:
            obs = torch.from_numpy(obs).permute(2, 0, 1).numpy()
        obs = self._greenscreen(obs)
        if channels_last:
            obs = torch.from_numpy(obs).permute(1, 2, 0).numpy()
        return obs

reason: import numpy,from numpy