<commit_before># -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Contains the transformation functions for getting from ICRS/HCRS to CIRS and
anything in between (currently that means GCRS)
"""

import numpy as np

from astropy import units as u
from astropy.coordinates.baseframe import frame_transform_graph
from astropy.coordinates.transformations import (
    FunctionTransformWithFiniteDifference,
    AffineTransform,
)
from astropy.coordinates.representation import (
    SphericalRepresentation,
    CartesianRepresentation,
    UnitSphericalRepresentation,
    CartesianDifferential,
)

from .icrs import ICRS
from .gcrs import GCRS
from .cirs import CIRS
from .hcrs import HCRS
from .equatorial import TETE
from .utils import aticq, atciqz

from ..erfa_astrom import erfa_astrom

# First the ICRS/CIRS related transforms
@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, ICRS, CIRS)
def icrs_to_cirs(icrs_coo, cirs_frame):
    # first set up the astrometry context for ICRS<->CIRS
    astrom = erfa_astrom.get().apci(cirs_frame)

    if icrs_coo.data.get_name() == 'unitspherical' or icrs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just do the infinite-distance/no parallax calculation
        srepr = icrs_coo.spherical
        cirs_ra, cirs_dec = atciqz(srepr.without_differentials(), astrom)

        newrep = UnitSphericalRepresentation(lat=u.Quantity(cirs_dec, u.radian, copy=False),
                                             lon=u.Quantity(cirs_ra, u.radian, copy=False),
                                             copy=False)
    else:
        # When there is a distance,  we first offset for parallax to get the
        # astrometric coordinate direction and *then* run the ERFA transform for
        # no parallax/PM. This ensures reversibility and is more sensible for
        # inside solar system objects
        astrom_eb = CartesianRepresentation(astrom['eb'], unit=u.au,
                                            xyz_axis=-1, copy=False)
        newcart = icrs_coo.cartesian - astrom_eb
        srepr = newcart.represent_as(SphericalRepresentation)
        cirs_ra, cirs_dec = atciqz(srepr.without_differentials(), astrom)

        newrep = SphericalRepresentation(lat=u.Quantity(cirs_dec, u.radian, copy=False),
                                         lon=u.Quantity(cirs_ra, u.radian, copy=False),
                                         distance=srepr.distance, copy=False)

    return cirs_frame.realize_frame(newrep)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, CIRS, ICRS)
def cirs_to_icrs(cirs_coo, icrs_frame):
    # set up the astrometry context for ICRS<->cirs and then convert to
    # astrometric coordinate direction
    astrom = erfa_astrom.get().apci(cirs_coo)
    srepr = cirs_coo.represent_as(SphericalRepresentation)
    i_ra, i_dec = aticq(srepr.without_differentials(), astrom)

    if cirs_coo.data.get_name() == 'unitspherical' or cirs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just use the coordinate direction to yield the
        # infinite-distance/no parallax answer
        newrep = UnitSphericalRepresentation(lat=u.Quantity(i_dec, u.radian, copy=False),
                                             lon=u.Quantity(i_ra, u.radian, copy=False),
                                             copy=False)
    else:
        # When there is a distance, apply the parallax/offset to the SSB as the
        # last step - ensures round-tripping with the icrs_to_cirs transform

        # the distance in intermedrep is *not* a real distance as it does not
        # include the offset back to the SSB
        intermedrep = SphericalRepresentation(lat=u.Quantity(i_dec, u.radian, copy=False),
                                              lon=u.Quantity(i_ra, u.radian, copy=False),
                                              distance=srepr.distance,
                                              copy=False)

        astrom_eb = CartesianRepresentation(astrom['eb'], unit=u.au,
                                            xyz_axis=-1, copy=False)
        newrep = intermedrep + astrom_eb

    return icrs_frame.realize_frame(newrep)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, CIRS, CIRS)
def cirs_to_cirs(from_coo, to_frame):
    if np.all(from_coo.obstime == to_frame.obstime):
        return to_frame.realize_frame(from_coo.data)
    else:
        # the CIRS<-> CIRS transform actually goes through ICRS.  This has a
        # subtle implication that a point in CIRS is uniquely determined
        # by the corresponding astrometric ICRS coordinate *at its
        # current time*.  This has some subtle implications in terms of GR, but
        # is sort of glossed over in the current scheme because we are dropping
        # distances anyway.
        return from_coo.transform_to(ICRS()).transform_to(to_frame)

# Now the GCRS-related transforms to/from ICRS

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, ICRS, GCRS)
def icrs_to_gcrs(icrs_coo, gcrs_frame):
    # first set up the astrometry context for ICRS<->GCRS.
    astrom = erfa_astrom.get().apcs(gcrs_frame)

    if icrs_coo.data.get_name() == 'unitspherical' or icrs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just do the infinite-distance/no parallax calculation
        srepr = icrs_coo.represent_as(SphericalRepresentation)
        gcrs_ra, gcrs_dec = atciqz(srepr.without_differentials(), astrom)

        newrep = UnitSphericalRepresentation(lat=u.Quantity(gcrs_dec, u.radian, copy=False),
                                             lon=u.Quantity(gcrs_ra, u.radian, copy=False),
                                             copy=False)
    else:
        # When there is a distance,  we first offset for parallax to get the
        # BCRS coordinate direction and *then* run the ERFA transform for no
        # parallax/PM. This ensures reversibility and is more sensible for
        # inside solar system objects
        astrom_eb = CartesianRepresentation(astrom['eb'], unit=u.au,
                                            xyz_axis=-1, copy=False)
        newcart = icrs_coo.cartesian - astrom_eb

        srepr = newcart.represent_as(SphericalRepresentation)
        gcrs_ra, gcrs_dec = atciqz(srepr.without_differentials(), astrom)

        newrep = SphericalRepresentation(lat=u.Quantity(gcrs_dec, u.radian, copy=False),
                                         lon=u.Quantity(gcrs_ra, u.radian, copy=False),
                                         distance=srepr.distance, copy=False)

    return gcrs_frame.realize_frame(newrep)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference,
                                 GCRS, ICRS)
def gcrs_to_icrs(gcrs_coo, icrs_frame):
    # set up the astrometry context for ICRS<->GCRS and then convert to BCRS
    # coordinate direction
    astrom = erfa_astrom.get().apcs(gcrs_coo)

    srepr = gcrs_coo.represent_as(SphericalRepresentation)
    i_ra, i_dec = aticq(srepr.without_differentials(), astrom)

    if gcrs_coo.data.get_name() == 'unitspherical' or gcrs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just use the coordinate direction to yield the
        # infinite-distance/no parallax answer
        newrep = UnitSphericalRepresentation(lat=u.Quantity(i_dec, u.radian, copy=False),
                                             lon=u.Quantity(i_ra, u.radian, copy=False),
                                             copy=False)
    else:
        # When there is a distance, apply the parallax/offset to the SSB as the
        # last step - ensures round-tripping with the icrs_to_gcrs transform

        # the distance in intermedrep is *not* a real distance as it does not
        # include the offset back to the SSB
        intermedrep = SphericalRepresentation(lat=u.Quantity(i_dec, u.radian, copy=False),
                                              lon=u.Quantity(i_ra, u.radian, copy=False),
                                              distance=srepr.distance,
                                              copy=False)

        astrom_eb = CartesianRepresentation(astrom['eb'], unit=u.au,
                                            xyz_axis=-1, copy=False)
        newrep = intermedrep + astrom_eb

    return icrs_frame.realize_frame(newrep)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, GCRS, GCRS)
def gcrs_to_gcrs(from_coo, to_frame):
    if (np.all(from_coo.obstime == to_frame.obstime) and
            np.all(from_coo.obsgeoloc == to_frame.obsgeoloc)):
        return to_frame.realize_frame(from_coo.data)
    else:
        # like CIRS, we do this self-transform via ICRS
        return from_coo.transform_to(ICRS()).transform_to(to_frame)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, GCRS, HCRS)
def gcrs_to_hcrs(gcrs_coo, hcrs_frame):

    if np.any(gcrs_coo.obstime != hcrs_frame.obstime):
        # if they GCRS obstime and HCRS obstime are not the same, we first
        # have to move to a GCRS where they are.
        frameattrs = gcrs_coo.get_frame_attr_names()
        frameattrs['obstime'] = hcrs_frame.obstime
        gcrs_coo = gcrs_coo.transform_to(GCRS(**frameattrs))

    # set up the astrometry context for ICRS<->GCRS and then convert to ICRS
    # coordinate direction
    astrom = erfa_astrom.get().apcs(gcrs_coo)
    srepr = gcrs_coo.represent_as(SphericalRepresentation)
    i_ra, i_dec = aticq(srepr.without_differentials(), astrom)

    # convert to Quantity objects
    i_ra = u.Quantity(i_ra, u.radian, copy=False)
    i_dec = u.Quantity(i_dec, u.radian, copy=False)
    if gcrs_coo.data.get_name() == 'unitspherical' or gcrs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just use the coordinate direction to yield the
        # infinite-distance/no parallax answer
        newrep = UnitSphericalRepresentation(lat=i_dec, lon=i_ra, copy=False)
    else:
        # When there is a distance, apply the parallax/offset to the
        # Heliocentre as the last step to ensure round-tripping with the
        # hcrs_to_gcrs transform

        # Note that the distance in intermedrep is *not* a real distance as it
        # does not include the offset back to the Heliocentre
        intermedrep = SphericalRepresentation(lat=i_dec, lon=i_ra,
                                              distance=srepr.distance,
                                              copy=False)

        # astrom['eh'] and astrom['em'] contain Sun to observer unit vector,
        # and distance, respectively. Shapes are (X) and (X,3), where (X) is the
        # shape resulting from broadcasting the shape of the times object
        # against the shape of the pv array.
        # broadcast em to eh and scale eh
        eh = astrom['eh'] * astrom['em'][..., np.newaxis]
        eh = CartesianRepresentation(eh, unit=u.au, xyz_axis=-1, copy=False)

        newrep = intermedrep.to_cartesian() + eh

    return hcrs_frame.realize_frame(newrep)

_NEED_ORIGIN_HINT = ("The input {0} coordinates do not have length units. This "
                     "probably means you created coordinates with lat/lon but "
                     "no distance.  Heliocentric<->ICRS transforms cannot "
                     "function in this case because there is an origin shift.")

@frame_transform_graph.transform(AffineTransform, HCRS, ICRS)
def hcrs_to_icrs(hcrs_coo, icrs_frame):
    # this is just an origin translation so without a distance it cannot go ahead
    if isinstance(hcrs_coo.data, UnitSphericalRepresentation):
        raise u.UnitsError(_NEED_ORIGIN_HINT.format(hcrs_coo.__class__.__name__))

    if hcrs_coo.data.differentials:
        from astropy.coordinates.solar_system import get_body_barycentric_posvel
        bary_sun_pos, bary_sun_vel = get_body_barycentric_posvel('sun',
                                                                 hcrs_coo.obstime)
        bary_sun_vel = bary_sun_vel.represent_as(CartesianDifferential)
        bary_sun_pos = bary_sun_pos.with_differentials(bary_sun_vel)

    else:
        from astropy.coordinates.solar_system import get_body_barycentric
        bary_sun_pos = get_body_barycentric('sun', hcrs_coo.obstime)
        bary_sun_vel = None

    return None, bary_sun_pos

@frame_transform_graph.transform(AffineTransform, ICRS, HCRS)
def icrs_to_hcrs(icrs_coo, hcrs_frame):
    # this is just an origin translation so without a distance it cannot go ahead
    if isinstance(icrs_coo.data, UnitSphericalRepresentation):
        raise u.UnitsError(_NEED_ORIGIN_HINT.format(icrs_coo.__class__.__name__))

    if icrs_coo.data.differentials:
        from astropy.coordinates.solar_system import get_body_barycentric_posvel
        bary_sun_pos, bary_sun_vel = get_body_barycentric_posvel('sun',
                                                                 hcrs_frame.obstime)
        bary_sun_pos = -bary_sun_pos
        bary_sun_vel = -bary_sun_vel.represent_as(CartesianDifferential)
        bary_sun_pos = bary_sun_pos.with_differentials(bary_sun_vel)

    else:
        from astropy.coordinates.solar_system import get_body_barycentric
        bary_sun_pos = -get_body_barycentric('sun', hcrs_frame.obstime)
        bary_sun_vel = None

    return None, bary_sun_pos

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, HCRS, HCRS)
def hcrs_to_hcrs(from_coo, to_frame):
    if np.all(from_coo.obstime == to_frame.obstime):
        return to_frame.realize_frame(from_coo.data)
    else:
        # like CIRS, we do this self-transform via ICRS
        return from_coo.transform_to(ICRS()).transform_to(to_frame)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, TETE, TETE)
def tete_to_tete(from_coo, to_frame):
    if (np.all(from_coo.obstime == to_frame.obstime) and
            np.all(from_coo.obsgeoloc == to_frame.obsgeoloc) and
            np.all(from_coo.obsgeovel == to_frame.obsgeovel)):
        return to_frame.realize_frame(from_coo.data)
    else:
        # can't do self-transform via ITRS due to infinite recursion
        # therefore we will go through ICRS, a la CIRS and GCRS
        return from_coo.transform_to(ICRS()).transform_to(to_frame)
<commit_msg>clearer comment<commit_after># -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Contains the transformation functions for getting from ICRS/HCRS to CIRS and
anything in between (currently that means GCRS)
"""

import numpy as np

from astropy import units as u
from astropy.coordinates.baseframe import frame_transform_graph
from astropy.coordinates.transformations import (
    FunctionTransformWithFiniteDifference,
    AffineTransform,
)
from astropy.coordinates.representation import (
    SphericalRepresentation,
    CartesianRepresentation,
    UnitSphericalRepresentation,
    CartesianDifferential,
)

from .icrs import ICRS
from .gcrs import GCRS
from .cirs import CIRS
from .hcrs import HCRS
from .equatorial import TETE
from .utils import aticq, atciqz

from ..erfa_astrom import erfa_astrom

# First the ICRS/CIRS related transforms
@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, ICRS, CIRS)
def icrs_to_cirs(icrs_coo, cirs_frame):
    # first set up the astrometry context for ICRS<->CIRS
    astrom = erfa_astrom.get().apci(cirs_frame)

    if icrs_coo.data.get_name() == 'unitspherical' or icrs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just do the infinite-distance/no parallax calculation
        srepr = icrs_coo.spherical
        cirs_ra, cirs_dec = atciqz(srepr.without_differentials(), astrom)

        newrep = UnitSphericalRepresentation(lat=u.Quantity(cirs_dec, u.radian, copy=False),
                                             lon=u.Quantity(cirs_ra, u.radian, copy=False),
                                             copy=False)
    else:
        # When there is a distance,  we first offset for parallax to get the
        # astrometric coordinate direction and *then* run the ERFA transform for
        # no parallax/PM. This ensures reversibility and is more sensible for
        # inside solar system objects
        astrom_eb = CartesianRepresentation(astrom['eb'], unit=u.au,
                                            xyz_axis=-1, copy=False)
        newcart = icrs_coo.cartesian - astrom_eb
        srepr = newcart.represent_as(SphericalRepresentation)
        cirs_ra, cirs_dec = atciqz(srepr.without_differentials(), astrom)

        newrep = SphericalRepresentation(lat=u.Quantity(cirs_dec, u.radian, copy=False),
                                         lon=u.Quantity(cirs_ra, u.radian, copy=False),
                                         distance=srepr.distance, copy=False)

    return cirs_frame.realize_frame(newrep)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, CIRS, ICRS)
def cirs_to_icrs(cirs_coo, icrs_frame):
    # set up the astrometry context for ICRS<->cirs and then convert to
    # astrometric coordinate direction
    astrom = erfa_astrom.get().apci(cirs_coo)
    srepr = cirs_coo.represent_as(SphericalRepresentation)
    i_ra, i_dec = aticq(srepr.without_differentials(), astrom)

    if cirs_coo.data.get_name() == 'unitspherical' or cirs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just use the coordinate direction to yield the
        # infinite-distance/no parallax answer
        newrep = UnitSphericalRepresentation(lat=u.Quantity(i_dec, u.radian, copy=False),
                                             lon=u.Quantity(i_ra, u.radian, copy=False),
                                             copy=False)
    else:
        # When there is a distance, apply the parallax/offset to the SSB as the
        # last step - ensures round-tripping with the icrs_to_cirs transform

        # the distance in intermedrep is *not* a real distance as it does not
        # include the offset back to the SSB
        intermedrep = SphericalRepresentation(lat=u.Quantity(i_dec, u.radian, copy=False),
                                              lon=u.Quantity(i_ra, u.radian, copy=False),
                                              distance=srepr.distance,
                                              copy=False)

        astrom_eb = CartesianRepresentation(astrom['eb'], unit=u.au,
                                            xyz_axis=-1, copy=False)
        newrep = intermedrep + astrom_eb

    return icrs_frame.realize_frame(newrep)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, CIRS, CIRS)
def cirs_to_cirs(from_coo, to_frame):
    if np.all(from_coo.obstime == to_frame.obstime):
        return to_frame.realize_frame(from_coo.data)
    else:
        # the CIRS<-> CIRS transform actually goes through ICRS.  This has a
        # subtle implication that a point in CIRS is uniquely determined
        # by the corresponding astrometric ICRS coordinate *at its
        # current time*.  This has some subtle implications in terms of GR, but
        # is sort of glossed over in the current scheme because we are dropping
        # distances anyway.
        return from_coo.transform_to(ICRS()).transform_to(to_frame)

# Now the GCRS-related transforms to/from ICRS

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, ICRS, GCRS)
def icrs_to_gcrs(icrs_coo, gcrs_frame):
    # first set up the astrometry context for ICRS<->GCRS.
    astrom = erfa_astrom.get().apcs(gcrs_frame)

    if icrs_coo.data.get_name() == 'unitspherical' or icrs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just do the infinite-distance/no parallax calculation
        srepr = icrs_coo.represent_as(SphericalRepresentation)
        gcrs_ra, gcrs_dec = atciqz(srepr.without_differentials(), astrom)

        newrep = UnitSphericalRepresentation(lat=u.Quantity(gcrs_dec, u.radian, copy=False),
                                             lon=u.Quantity(gcrs_ra, u.radian, copy=False),
                                             copy=False)
    else:
        # When there is a distance,  we first offset for parallax to get the
        # BCRS coordinate direction and *then* run the ERFA transform for no
        # parallax/PM. This ensures reversibility and is more sensible for
        # inside solar system objects
        astrom_eb = CartesianRepresentation(astrom['eb'], unit=u.au,
                                            xyz_axis=-1, copy=False)
        newcart = icrs_coo.cartesian - astrom_eb

        srepr = newcart.represent_as(SphericalRepresentation)
        gcrs_ra, gcrs_dec = atciqz(srepr.without_differentials(), astrom)

        newrep = SphericalRepresentation(lat=u.Quantity(gcrs_dec, u.radian, copy=False),
                                         lon=u.Quantity(gcrs_ra, u.radian, copy=False),
                                         distance=srepr.distance, copy=False)

    return gcrs_frame.realize_frame(newrep)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference,
                                 GCRS, ICRS)
def gcrs_to_icrs(gcrs_coo, icrs_frame):
    # set up the astrometry context for ICRS<->GCRS and then convert to BCRS
    # coordinate direction
    astrom = erfa_astrom.get().apcs(gcrs_coo)

    srepr = gcrs_coo.represent_as(SphericalRepresentation)
    i_ra, i_dec = aticq(srepr.without_differentials(), astrom)

    if gcrs_coo.data.get_name() == 'unitspherical' or gcrs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just use the coordinate direction to yield the
        # infinite-distance/no parallax answer
        newrep = UnitSphericalRepresentation(lat=u.Quantity(i_dec, u.radian, copy=False),
                                             lon=u.Quantity(i_ra, u.radian, copy=False),
                                             copy=False)
    else:
        # When there is a distance, apply the parallax/offset to the SSB as the
        # last step - ensures round-tripping with the icrs_to_gcrs transform

        # the distance in intermedrep is *not* a real distance as it does not
        # include the offset back to the SSB
        intermedrep = SphericalRepresentation(lat=u.Quantity(i_dec, u.radian, copy=False),
                                              lon=u.Quantity(i_ra, u.radian, copy=False),
                                              distance=srepr.distance,
                                              copy=False)

        astrom_eb = CartesianRepresentation(astrom['eb'], unit=u.au,
                                            xyz_axis=-1, copy=False)
        newrep = intermedrep + astrom_eb

    return icrs_frame.realize_frame(newrep)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, GCRS, GCRS)
def gcrs_to_gcrs(from_coo, to_frame):
    if (np.all(from_coo.obstime == to_frame.obstime) and
            np.all(from_coo.obsgeoloc == to_frame.obsgeoloc)):
        return to_frame.realize_frame(from_coo.data)
    else:
        # like CIRS, we do this self-transform via ICRS
        return from_coo.transform_to(ICRS()).transform_to(to_frame)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, GCRS, HCRS)
def gcrs_to_hcrs(gcrs_coo, hcrs_frame):

    if np.any(gcrs_coo.obstime != hcrs_frame.obstime):
        # if they GCRS obstime and HCRS obstime are not the same, we first
        # have to move to a GCRS where they are.
        frameattrs = gcrs_coo.get_frame_attr_names()
        frameattrs['obstime'] = hcrs_frame.obstime
        gcrs_coo = gcrs_coo.transform_to(GCRS(**frameattrs))

    # set up the astrometry context for ICRS<->GCRS and then convert to ICRS
    # coordinate direction
    astrom = erfa_astrom.get().apcs(gcrs_coo)
    srepr = gcrs_coo.represent_as(SphericalRepresentation)
    i_ra, i_dec = aticq(srepr.without_differentials(), astrom)

    # convert to Quantity objects
    i_ra = u.Quantity(i_ra, u.radian, copy=False)
    i_dec = u.Quantity(i_dec, u.radian, copy=False)
    if gcrs_coo.data.get_name() == 'unitspherical' or gcrs_coo.data.to_cartesian().x.unit == u.one:
        # if no distance, just use the coordinate direction to yield the
        # infinite-distance/no parallax answer
        newrep = UnitSphericalRepresentation(lat=i_dec, lon=i_ra, copy=False)
    else:
        # When there is a distance, apply the parallax/offset to the
        # Heliocentre as the last step to ensure round-tripping with the
        # hcrs_to_gcrs transform

        # Note that the distance in intermedrep is *not* a real distance as it
        # does not include the offset back to the Heliocentre
        intermedrep = SphericalRepresentation(lat=i_dec, lon=i_ra,
                                              distance=srepr.distance,
                                              copy=False)

        # astrom['eh'] and astrom['em'] contain Sun to observer unit vector,
        # and distance, respectively. Shapes are (X) and (X,3), where (X) is the
        # shape resulting from broadcasting the shape of the times object
        # against the shape of the pv array.
        # broadcast em to eh and scale eh
        eh = astrom['eh'] * astrom['em'][..., np.newaxis]
        eh = CartesianRepresentation(eh, unit=u.au, xyz_axis=-1, copy=False)

        newrep = intermedrep.to_cartesian() + eh

    return hcrs_frame.realize_frame(newrep)

_NEED_ORIGIN_HINT = ("The input {0} coordinates do not have length units. This "
                     "probably means you created coordinates with lat/lon but "
                     "no distance.  Heliocentric<->ICRS transforms cannot "
                     "function in this case because there is an origin shift.")

@frame_transform_graph.transform(AffineTransform, HCRS, ICRS)
def hcrs_to_icrs(hcrs_coo, icrs_frame):
    # this is just an origin translation so without a distance it cannot go ahead
    if isinstance(hcrs_coo.data, UnitSphericalRepresentation):
        raise u.UnitsError(_NEED_ORIGIN_HINT.format(hcrs_coo.__class__.__name__))

    if hcrs_coo.data.differentials:
        from astropy.coordinates.solar_system import get_body_barycentric_posvel
        bary_sun_pos, bary_sun_vel = get_body_barycentric_posvel('sun',
                                                                 hcrs_coo.obstime)
        bary_sun_vel = bary_sun_vel.represent_as(CartesianDifferential)
        bary_sun_pos = bary_sun_pos.with_differentials(bary_sun_vel)

    else:
        from astropy.coordinates.solar_system import get_body_barycentric
        bary_sun_pos = get_body_barycentric('sun', hcrs_coo.obstime)
        bary_sun_vel = None

    return None, bary_sun_pos

@frame_transform_graph.transform(AffineTransform, ICRS, HCRS)
def icrs_to_hcrs(icrs_coo, hcrs_frame):
    # this is just an origin translation so without a distance it cannot go ahead
    if isinstance(icrs_coo.data, UnitSphericalRepresentation):
        raise u.UnitsError(_NEED_ORIGIN_HINT.format(icrs_coo.__class__.__name__))

    if icrs_coo.data.differentials:
        from astropy.coordinates.solar_system import get_body_barycentric_posvel
        bary_sun_pos, bary_sun_vel = get_body_barycentric_posvel('sun',
                                                                 hcrs_frame.obstime)
        bary_sun_pos = -bary_sun_pos
        bary_sun_vel = -bary_sun_vel.represent_as(CartesianDifferential)
        bary_sun_pos = bary_sun_pos.with_differentials(bary_sun_vel)

    else:
        from astropy.coordinates.solar_system import get_body_barycentric
        bary_sun_pos = -get_body_barycentric('sun', hcrs_frame.obstime)
        bary_sun_vel = None

    return None, bary_sun_pos

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, HCRS, HCRS)
def hcrs_to_hcrs(from_coo, to_frame):
    if np.all(from_coo.obstime == to_frame.obstime):
        return to_frame.realize_frame(from_coo.data)
    else:
        # like CIRS, we do this self-transform via ICRS
        return from_coo.transform_to(ICRS()).transform_to(to_frame)

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, TETE, TETE)
def tete_to_tete(from_coo, to_frame):
    if (np.all(from_coo.obstime == to_frame.obstime) and
            np.all(from_coo.obsgeoloc == to_frame.obsgeoloc) and
            np.all(from_coo.obsgeovel == to_frame.obsgeovel)):
        return to_frame.realize_frame(from_coo.data)
    else:
        # We perform the self-transform through ICRS, since any change in obstime
        # will require handling the change in pos/vel of the Earth and it's impact
        # on aberration.
        return from_coo.transform_to(ICRS()).transform_to(to_frame)
<|endoftext|>


<commit_before>#   Copyright (c) 2018 PaddlePaddle 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.

from __future__ import print_function

import os
import signal
import time
import unittest
from multiprocessing import Process

import numpy as np
import paddle.fluid as fluid
import paddle.fluid.core as core
from paddle.fluid.op import Operator
from paddle.fluid.framework import Program, program_guard

def run_pserver(pserver_id, use_cuda, sync_mode):
    scope = fluid.core.Scope()
    program = Program()
    with fluid.scope_guard(scope):
        with program_guard(program, startup_program=Program()):
            # create table parameter in scope
            place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
            # create and initialize Param Variable
            param = scope.var('table').get_tensor()

            param_array = np.ones((5, 8)).astype("float32")
            for i in range(len(param_array)):
                param_array[i] *= param_array[i] * i + pserver_id * 10 + 1
            param.set(param_array, place)

            optimize_block = program._create_block(program.global_block().idx)
            program.global_block().append_op(
                type="listen_and_serv",
                inputs={'X': []},
                outputs={},
                attrs={
                    "optimize_blocks": [optimize_block],
                    "endpoint": '127.0.0.1:0',
                    "Fanin": 1,
                    "sync_mode": True,
                    "grad_to_block_id": []
                })

            exe = fluid.Executor(place)
            exe.run(program)

class TestListenAndServOp(unittest.TestCase):
    def setUp(self):
        self.ps_timeout = 5

    def _start_pserver(self, pserver_id, use_cuda, sync_mode, pserver_func):
        p = Process(target=pserver_func, args=(pserver_id, use_cuda, sync_mode))
        p.daemon = True
        p.start()
        return p

    def _wait_ps_ready(self, pid):
        start_left_time = self.ps_timeout
        sleep_time = 0.5
        while True:
            assert start_left_time >= 0, "wait ps ready failed"
            time.sleep(sleep_time)
            try:
                # the listen_and_serv_op would touch a file which contains the listen port
                # on the /tmp directory until it was ready to process all the RPC call.
                os.stat("/tmp/paddle.%d.port" % pid)
                return
            except os.error:
                start_left_time -= sleep_time

    def _get_pserver_port(self, pid):
        with open("/tmp/paddle.%d.port" % pid, 'r') as f:
            port = int(f.read().strip())
        return port

    def _run_nce_op_two_pserver(self, place, port0, port1):
        scope = fluid.core.Scope()
        program = Program()
        with fluid.scope_guard(scope):
            with program_guard(program, startup_program=Program()):
                x = scope.var('Input').get_tensor()
                x_array = np.random.random((4, 8)).astype("float32") * 2
                x.set(x_array, place)
                # create and initialize Param Variable
                param = scope.var('Weight').get_tensor()
                param_array = np.zeros((5, 8)).astype("float32") * 2
                param.set(param_array, place)

                bias = scope.var('Bias').get_tensor()
                bias_array = np.random.random((5, 1)).astype("float32")
                bias.set(bias_array, place)

                sample_w = scope.var('SampleWeight').get_tensor()
                sample_weight = np.random.random((4, 1)).astype("float32")
                sample_w.set(sample_weight, place)

                label = scope.var('Label').get_tensor()
                label_array = np.array([0, 1, 4, 5])
                label.set(label_array, place)

                cost = scope.var('Cost').get_tensor()
                cost_w = np.zeros((4, 1)).astype("float32")
                cost.set(cost_w, place)

                sample_l = scope.var('SampleLogits').get_tensor()
                sample_l_w = np.zeros((4, 3)).astype("float32")
                sample_l.set(sample_l_w, place)

                sample_la = scope.var('SampleLabels').get_tensor()
                sample_la_w = np.zeros((4, 3)).astype("float32")
                sample_la.set(sample_la_w, place)

                emaps = ['127.0.0.1:' + str(port0), '127.0.0.1:' + str(port1)]
                table_names = ['table', 'table']
                height_sections = [2, 3]

                # create and run nce operator
                nce_op = Operator(
                    "nce",
                    Input='Input',
                    Weight='Weight',
                    Label='Label',
                    Bias='Bias',
                    Cost='Cost',
                    SampleLogits='SampleLogits',
                    SampleLabels='SampleLabels',
                    num_total_classes=5,
                    num_neg_samples=2,
                    custom_neg_classes=list(range(2)),
                    sampler=0,
                    seed=1,
                    is_sparse=True,
                    remote_prefetch=True,
                    epmap=emaps,
                    table_names=table_names,
                    height_sections=height_sections)

                nce_op.run(scope, place)

                # get and compare result
                o_cost = np.array(cost_w)
                o_logits = np.array(sample_l)
                o_labels = np.array(sample_la)

    def test_nce_op_remote(self):
        os.environ['PADDLE_ENABLE_REMOTE_PREFETCH'] = "1"
        # run pserver on CPU in sync mode
        p0 = self._start_pserver(0, False, True, run_pserver)
        self._wait_ps_ready(p0.pid)
        port0 = self._get_pserver_port(p0.pid)

        p1 = self._start_pserver(1, False, True, run_pserver)
        self._wait_ps_ready(p1.pid)
        port1 = self._get_pserver_port(p1.pid)

        places = [core.CPUPlace()]
        if core.is_compiled_with_cuda():
            places.append(core.CUDAPlace(0))

        for place in places:
            self._run_nce_op_two_pserver(place, port0, port1)

        # raise SIGTERM to pserver
        os.kill(p0.pid, signal.SIGINT)
        p0.join()
        os.kill(p1.pid, signal.SIGINT)
        p1.join()

if __name__ == '__main__':
    unittest.main()
<commit_msg>add nce remote ut, test=develop<commit_after>#   Copyright (c) 2018 PaddlePaddle 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.

from __future__ import print_function

import os
import signal
import time
import unittest
from multiprocessing import Process

import numpy as np
import paddle.fluid as fluid
import paddle.fluid.core as core
from paddle.fluid.op import Operator
from paddle.fluid.framework import Program, program_guard

def nce(input, weight, bias, sample_weight, labels, num_classes,
        num_sample_class):
    samples = []
    sample_labels = []
    batch_size = input.shape[0]
    num_true_class = labels.shape[1]
    for i in range(batch_size):
        w = 1 if sample_weight is None else sample_weight[i]
        for label in labels[i]:
            samples.append((i, label, True, w))
            sample_labels.append(label)
        for num in range(num_sample_class):
            samples.append((i, num, False, w))
            sample_labels.append(num)
    # forward bias
    sample_out = np.zeros(len(samples)).astype(np.float32)
    if bias is not None:
        for i in range(len(samples)):
            sample_out[i] = bias[samples[i][1]]
    # forward weight
    for i in range(len(samples)):
        sample_out[i] += np.dot(input[samples[i][0]], weight[samples[i][1]])

    # forward activation
    sample_out = 1.0 / (1.0 + np.exp(-sample_out))
    # forward cost
    out = np.zeros(batch_size).astype(np.float32)
    b = 1.0 / num_classes * num_sample_class

    for i in range(len(samples)):
        o = sample_out[i]
        cost = -np.log(o / (o + b)) if samples[i][2] else -np.log(b / (o + b))
        out[samples[i][0]] += cost * samples[i][3]
    return (out[:, np.newaxis], np.array(sample_out).reshape(
        batch_size, num_sample_class + num_true_class),
            np.array(sample_labels).reshape(batch_size,
                                            num_sample_class + num_true_class))

def run_pserver(pserver_id, use_cuda, sync_mode):
    scope = fluid.core.Scope()
    program = Program()
    with fluid.scope_guard(scope):
        with program_guard(program, startup_program=Program()):
            # create table parameter in scope
            place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
            # create and initialize Param Variable
            param = scope.var('table').get_tensor()

            param_array = np.ones((5, 8)).astype("float32")
            for i in range(len(param_array)):
                param_array[i] *= param_array[i] * i + pserver_id * 10 + 1
            param.set(param_array, place)

            optimize_block = program._create_block(program.global_block().idx)
            program.global_block().append_op(
                type="listen_and_serv",
                inputs={'X': []},
                outputs={},
                attrs={
                    "optimize_blocks": [optimize_block],
                    "endpoint": '127.0.0.1:0',
                    "Fanin": 1,
                    "sync_mode": True,
                    "grad_to_block_id": []
                })

            exe = fluid.Executor(place)
            exe.run(program)

class TestListenAndServOp(unittest.TestCase):
    def setUp(self):
        self.ps_timeout = 5

    def _start_pserver(self, pserver_id, use_cuda, sync_mode, pserver_func):
        p = Process(target=pserver_func, args=(pserver_id, use_cuda, sync_mode))
        p.daemon = True
        p.start()
        return p

    def _wait_ps_ready(self, pid):
        start_left_time = self.ps_timeout
        sleep_time = 0.5
        while True:
            assert start_left_time >= 0, "wait ps ready failed"
            time.sleep(sleep_time)
            try:
                # the listen_and_serv_op would touch a file which contains the listen port
                # on the /tmp directory until it was ready to process all the RPC call.
                os.stat("/tmp/paddle.%d.port" % pid)
                return
            except os.error:
                start_left_time -= sleep_time

    def _get_pserver_port(self, pid):
        with open("/tmp/paddle.%d.port" % pid, 'r') as f:
            port = int(f.read().strip())
        return port

    def _run_nce_op_two_pserver(self, place, port0, port1):
        scope = fluid.core.Scope()
        program = Program()
        with fluid.scope_guard(scope):
            with program_guard(program, startup_program=Program()):
                x = scope.var('Input').get_tensor()
                x_array = np.random.random((4, 8)).astype("float32")
                x.set(x_array, place)
                # create and initialize Param Variable
                param = scope.var('Weight').get_tensor()
                param_array = np.zeros((5, 8)).astype("float32")
                param.set(param_array, place)

                bias = scope.var('Bias').get_tensor()
                bias_array = np.random.random((5, 1)).astype("float32")
                bias.set(bias_array, place)

                sample_w = scope.var('SampleWeight').get_tensor()
                sample_weight = np.random.random((4, 1)).astype("float32")
                sample_w.set(sample_weight, place)

                label = scope.var('Label').get_tensor()
                label_array = np.array([[0], [1], [4], [3]])
                label.set(label_array, place)

                cost = scope.var('Cost').get_tensor()
                cost_w = np.zeros((4, 1)).astype("float32")
                cost.set(cost_w, place)

                sample_l = scope.var('SampleLogits').get_tensor()
                sample_l_w = np.zeros((4, 3)).astype("float32")
                sample_l.set(sample_l_w, place)

                sample_la = scope.var('SampleLabels').get_tensor()
                sample_la_w = np.zeros((4, 3)).astype("int")
                sample_la.set(sample_la_w, place)

                emaps = ['127.0.0.1:' + str(port0), '127.0.0.1:' + str(port1)]
                table_names = ['table', 'table']
                height_sections = [2, 3]

                # create and run nce operator
                nce_op = Operator(
                    "nce",
                    Input='Input',
                    Weight='Weight',
                    Label='Label',
                    Bias='Bias',
                    Cost='Cost',
                    SampleLogits='SampleLogits',
                    SampleLabels='SampleLabels',
                    SampleWeight='SampleWeight',
                    num_total_classes=5,
                    num_neg_samples=2,
                    custom_neg_classes=list(range(2)),
                    sampler=0,
                    seed=0,
                    is_sparse=True,
                    remote_prefetch=True,
                    epmap=emaps,
                    table_names=table_names,
                    height_sections=height_sections)

                nce_op.run(scope, place)

                # get and compare result
                o_cost = np.array(scope.var('Cost').get_tensor())
                o_logits = np.array(scope.var('SampleLogits').get_tensor())
                o_labels = np.array(scope.var('SampleLabels').get_tensor())

                param_array = np.ones((5, 8)).astype("float32")
                for i in range(2):
                    param_array[i] *= param_array[i] * i + 0 * 10 + 1
                for i in range(2, 5):
                    param_array[i] *= param_array[i] * i + 1 * 10 + 1
                out = nce(x_array, param_array, bias_array, sample_weight,
                          label_array, 5, 2)

                self.assertAlmostEqual(o_cost.all(), out[0].all(), delta=1e-6)
                self.assertAlmostEqual(o_logits.all(), out[1].all(), delta=1e-6)
                self.assertAlmostEqual(o_labels.all(), out[2].all(), delta=1e-6)

    def test_nce_op_remote(self):
        os.environ['PADDLE_ENABLE_REMOTE_PREFETCH'] = "1"
        # run pserver on CPU in sync mode
        p0 = self._start_pserver(0, False, True, run_pserver)
        self._wait_ps_ready(p0.pid)
        port0 = self._get_pserver_port(p0.pid)

        p1 = self._start_pserver(1, False, True, run_pserver)
        self._wait_ps_ready(p1.pid)
        port1 = self._get_pserver_port(p1.pid)

        places = [core.CPUPlace()]
        if core.is_compiled_with_cuda():
            places.append(core.CUDAPlace(0))

        for place in places:
            self._run_nce_op_two_pserver(place, port0, port1)

        # raise SIGTERM to pserver
        os.kill(p0.pid, signal.SIGINT)
        p0.join()
        os.kill(p1.pid, signal.SIGINT)
        p1.join()

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


<commit_before>"""
Class to skim the PDB for similar unit cells
@author: hlasimpk
"""

# Imports
from Bio.PDB import PDBList
import os
import numpy as np
from collections import namedtuple
import cctbx.crystal
import cctbx.uctbx
import cPickle
import gzip
import logging

from simbad.constants import SIMBAD_LATTICE_DB

logger = logging.getLogger(__name__)

class Lattice_search(object):
    def __init__(self, optd=None):

        self.unit_cell = None
        self.space_group = None
        self.database = []
        self.niggli_cell = None
        self.results = []
        self.pickle_file = SIMBAD_LATTICE_DB
        self.total_files = 0

        if optd:
            self.init(optd)

        if self.unit_cell and self.space_group:
            self.calculate_niggli_cell()
            logger.info("Niggli cell calculated as: {0}".format(self.niggli_cell))

        if not self.database:
            self.load_database()

        if self.database:
            Lattice_parameter_scores = namedtuple("Lattice_parameter_scores", ['PDB_code',
                                                                               'unit_cell',
                                                                               'Penalty_score',
                                                                               'Length_penalty',
                                                                               'Angle_penalty'])

            item_ct = 0
            lattice_pars = self.database[1]
            for db_cell in lattice_pars:
                pdb_code = self.database[0][item_ct]
                try:
                    # Calculate the tolerances for the input niggli cell once
                    self.tolerance_a = ((self.niggli_cell[0] / 100) * 5)
                    self.tolerance_b = ((self.niggli_cell[1] / 100) * 5)
                    self.tolerance_c = ((self.niggli_cell[2] / 100) * 5)
                    self.tolerance_alpha = ((self.niggli_cell[3] / 100) * 5)
                    self.tolerance_beta = ((self.niggli_cell[4] / 100) * 5)
                    self.tolerance_gamma = ((self.niggli_cell[5] / 100) * 5)
                    within_tol = self.cell_tolerance(db_cell)
                except ValueError:
                    within_tol = False

                if within_tol:
                    total_penalty, length_penalty, angle_penalty = self.calculate_penalty(db_cell)
                    a, b, c, alpha, beta, gamma = db_cell
                    self.results.append(Lattice_parameter_scores(PDB_code=pdb_code,
                                                                 unit_cell='{0},{1},{2},{3},{4},{5}'.format(a, b, c,
                                                                                alpha, beta,
                                                                                gamma),
                                                                 Penalty_score="%08.5f" % float(abs(total_penalty)),
                                                                 Length_penalty="%08.5f" % float(abs(length_penalty)),
                                                                 Angle_penalty="%08.5f" % float(abs(angle_penalty))))
                item_ct += 1

        if self.results:
            self.rank_scores()
            self.return_results()
            if not optd['PDB']:
                logger.info("Downloading structures of the top 20 results")
                self.download_pdbs(optd)
            else:
                logger.info("Obtaining structures for the top 20 results from: {0}".format(optd['PDB']))
                self.copy_pdbs(optd)

        return

    def init(self, optd):
        '''Set input arguments as class variables'''

        if 'cell_parameters' in optd.keys() and optd['cell_parameters']:
            self.unit_cell = optd['cell_parameters']
        else:
            msg = "Cell parameters not specified"
            logger.critical(msg)
            raise RuntimeError(msg)

        if 'space_group' in optd.keys() and optd['space_group']:
            self.space_group = optd['space_group']
        else:
            msg = "Space group not specified"
            logger.critical(msg)
            raise RuntimeError(msg)

        return

    def calculate_niggli_cell(self):
        '''Function to calculate the niggli cell'''

        # Check for known anomalies, may need to add to this
        if self.space_group == "B2":
            self.space_group = "B112"
        elif self.space_group == "C1211":
            self.space_group = "C2"
        elif self.space_group == "P21212A":
            self.space_group = "P212121"
        elif self.space_group == "R3":
            self.space_group = "R3:R"
        elif self.space_group == "C4212":
            self.space_group = "P422"

        unit_cell = cctbx.uctbx.unit_cell(self.unit_cell)

        xs = cctbx.crystal.symmetry(
            unit_cell=unit_cell,
            space_group=self.space_group
        )
        niggli_cell = xs.change_basis(xs.change_of_basis_op_to_niggli_cell()).unit_cell()
        self.niggli_cell = np.array(eval(str(niggli_cell)))
        return

    def load_database(self):
        '''Function to load the database'''

        self.database = cPickle.load(open(self.pickle_file))
        self.total_files = len(self.database[0])
        logger.info("  - database loaded from pickle file")

        return

    def calculate_penalty(self, ref_niggli_cell):
        '''Calculate the linear cell variation between unit cells'''
        # Try needed to account for ref niggli cells that didn't work
        try:
            length_penalty, angle_penalty = self.penalty(ref_niggli_cell)

            total_penalty = length_penalty + angle_penalty

            return total_penalty, length_penalty, angle_penalty
        except ValueError:
            return None

    def cell_tolerance(self, ref_niggli_cell):
        '''Function to remove cells outside certain tolerances from being considered'''

        # Calculate if the 6 lattice parameters are within tolerances
        a = np.allclose(self.niggli_cell[0], ref_niggli_cell[0], atol=self.tolerance_a)
        b = np.allclose(self.niggli_cell[1], ref_niggli_cell[1], atol=self.tolerance_b)
        c = np.allclose(self.niggli_cell[2], ref_niggli_cell[2], atol=self.tolerance_c)
        alpha = np.allclose(self.niggli_cell[3], ref_niggli_cell[3], atol=self.tolerance_alpha)
        beta = np.allclose(self.niggli_cell[4], ref_niggli_cell[4], atol=self.tolerance_beta)
        gamma = np.allclose(self.niggli_cell[5], ref_niggli_cell[5], atol=self.tolerance_gamma)

        within_tol = False
        if a and b and c and alpha and beta and gamma:
            within_tol = True

        return within_tol

    def penalty(self, input_cell):
        '''Function to add penalty score depending on difference between unit cells'''
        # Split the experimental niggli cell into into component parts
        a_e, b_e, c_e, alpha_e, beta_e, gamma_e = self.niggli_cell

        # Split the input niggli cell into into component parts
        a_r, b_r, c_r, alpha_r, beta_r, gamma_r = input_cell

        delta_a = abs(a_e - a_r)
        delta_b = abs(b_e - b_r)
        delta_c = abs(c_e - c_r)
        delta_alpha = abs(alpha_e - alpha_r)
        delta_beta = abs(beta_e - beta_r)
        delta_gamma = abs(gamma_e - gamma_r)

        lengths_penalty_score = delta_a + delta_b + delta_c
        angles_penalty_score = delta_alpha + delta_beta + delta_gamma
        return lengths_penalty_score, angles_penalty_score

    def rank_scores(self):
        '''Rank the scores in self.results by Penalty score'''
        self.results = sorted(self.results, key=lambda x: float(x.Penalty_score), reverse=False)
        return

    def return_results(self):
        '''Log the results from the lattice parameter search'''
        # Output headers for the log file table
        logger.info("The lattice parameter search found the following structures:")
        logger.info(
            "PDB_CODE |   A   |   B   |   C   | ALPHA |  BETA | GAMMA | LENGTH_PENALTY | ANGLE_PENALTY | TOTAL_PENALTY ")

        count = 0
        for result in self.results:
            # Create a CSV for reading later
            with open('lattice.csv', 'a') as f:
                if count < 20:
                    f.write("{0},{1},{2},{3},{4}\n".format(result.PDB_code,
                                                           result.unit_cell,
                                                           result.Length_penalty,
                                                           result.Angle_penalty,
                                                           result.Penalty_score))
                    count += 1

            # Output a table to the log file
            uc = result.unit_cell.split(',')
            logger.info(
                '  {0}   | {1} | {2} | {3} | {4} | {5} | {6} |    {7}    |    {8}    |    {9}'.format(result.PDB_code,
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[0]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[1]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[2]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[3]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[4]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[5]))),
                                                                                result.Length_penalty,
                                                                                result.Angle_penalty,
                                                                                result.Penalty_score, ))

        return

    def copy_pdbs(self, optd):
        """Copy across the top 20 PDB files identified by lattice parameter search from a local download of the PDB"""

        count = 0
        for result in self.results:
            if count <= 20:
                with gzip.open(
                        os.path.join(optd['PDB'], '{0}', 'pdb{1}.ent.gz'.format(result.PDB_code[1:3], result.PDB_code)),
                        'rb') as f:
                    file_content = f.read()
                    with open(os.path.join(optd['work_dir'], 'lattice_input_models', '{0}.pdb'.format(result.PDB_code),
                                           'w')) as o:
                        o.write(file_content)
                count += 1
        return

    def download_pdbs(self, optd):
        '''Download the top 20 results directly from the PDB'''
        pdbl = PDBList()
        count = 0
        for result in self.results:
            if count < 20:
                # Download PDB file
                pdbl.retrieve_pdb_file(result.PDB_code, pdir=os.path.join(optd['work_dir'], 'lattice_input_models'))

                # Rename the PDB file as appropriate
                os.rename(
                    os.path.join(optd['work_dir'], 'lattice_input_models', 'pdb{0}.ent'.format(result.PDB_code.lower())),
                    os.path.join(optd['work_dir'], 'lattice_input_models', '{0}.pdb'.format(result.PDB_code)))
                count += 1

        return

def return_result_list():
    result_list = []
    with open('lattice.csv', 'r') as f:
        for line in f:
            PDB_code = line.split(',')[0]
            result_list.append(PDB_code)

    return result_list<commit_msg>removed logger line<commit_after>"""
Class to skim the PDB for similar unit cells
@author: hlasimpk
"""

# Imports
from Bio.PDB import PDBList
import os
import numpy as np
from collections import namedtuple
import cctbx.crystal
import cctbx.uctbx
import cPickle
import gzip
import logging

from simbad.constants import SIMBAD_LATTICE_DB

logger = logging.getLogger(__name__)

class Lattice_search(object):
    def __init__(self, optd=None):

        self.unit_cell = None
        self.space_group = None
        self.database = []
        self.niggli_cell = None
        self.results = []
        self.pickle_file = SIMBAD_LATTICE_DB
        self.total_files = 0

        if optd:
            self.init(optd)

        if self.unit_cell and self.space_group:
            self.calculate_niggli_cell()
            logger.info("Niggli cell calculated as: {0}".format(self.niggli_cell))

        if not self.database:
            self.load_database()

        if self.database:
            Lattice_parameter_scores = namedtuple("Lattice_parameter_scores", ['PDB_code',
                                                                               'unit_cell',
                                                                               'Penalty_score',
                                                                               'Length_penalty',
                                                                               'Angle_penalty'])

            item_ct = 0
            lattice_pars = self.database[1]
            for db_cell in lattice_pars:
                pdb_code = self.database[0][item_ct]
                try:
                    # Calculate the tolerances for the input niggli cell once
                    self.tolerance_a = ((self.niggli_cell[0] / 100) * 5)
                    self.tolerance_b = ((self.niggli_cell[1] / 100) * 5)
                    self.tolerance_c = ((self.niggli_cell[2] / 100) * 5)
                    self.tolerance_alpha = ((self.niggli_cell[3] / 100) * 5)
                    self.tolerance_beta = ((self.niggli_cell[4] / 100) * 5)
                    self.tolerance_gamma = ((self.niggli_cell[5] / 100) * 5)
                    within_tol = self.cell_tolerance(db_cell)
                except ValueError:
                    within_tol = False

                if within_tol:
                    total_penalty, length_penalty, angle_penalty = self.calculate_penalty(db_cell)
                    a, b, c, alpha, beta, gamma = db_cell
                    self.results.append(Lattice_parameter_scores(PDB_code=pdb_code,
                                                                 unit_cell='{0},{1},{2},{3},{4},{5}'.format(a, b, c,
                                                                                alpha, beta,
                                                                                gamma),
                                                                 Penalty_score="%08.5f" % float(abs(total_penalty)),
                                                                 Length_penalty="%08.5f" % float(abs(length_penalty)),
                                                                 Angle_penalty="%08.5f" % float(abs(angle_penalty))))
                item_ct += 1

        if self.results:
            self.rank_scores()
            self.return_results()
            if not optd['PDB']:
                logger.info("Downloading structures of the top 20 results")
                self.download_pdbs(optd)
            else:
                logger.info("Obtaining structures for the top 20 results from: {0}".format(optd['PDB']))
                self.copy_pdbs(optd)

        return

    def init(self, optd):
        '''Set input arguments as class variables'''

        if 'cell_parameters' in optd.keys() and optd['cell_parameters']:
            self.unit_cell = optd['cell_parameters']
        else:
            msg = "Cell parameters not specified"
            logger.critical(msg)
            raise RuntimeError(msg)

        if 'space_group' in optd.keys() and optd['space_group']:
            self.space_group = optd['space_group']
        else:
            msg = "Space group not specified"
            logger.critical(msg)
            raise RuntimeError(msg)

        return

    def calculate_niggli_cell(self):
        '''Function to calculate the niggli cell'''

        # Check for known anomalies, may need to add to this
        if self.space_group == "B2":
            self.space_group = "B112"
        elif self.space_group == "C1211":
            self.space_group = "C2"
        elif self.space_group == "P21212A":
            self.space_group = "P212121"
        elif self.space_group == "R3":
            self.space_group = "R3:R"
        elif self.space_group == "C4212":
            self.space_group = "P422"

        unit_cell = cctbx.uctbx.unit_cell(self.unit_cell)

        xs = cctbx.crystal.symmetry(
            unit_cell=unit_cell,
            space_group=self.space_group
        )
        niggli_cell = xs.change_basis(xs.change_of_basis_op_to_niggli_cell()).unit_cell()
        self.niggli_cell = np.array(eval(str(niggli_cell)))
        return

    def load_database(self):
        '''Function to load the database'''

        self.database = cPickle.load(open(self.pickle_file))
        self.total_files = len(self.database[0])
        logger.info("  - database loaded from pickle file")

        return

    def calculate_penalty(self, ref_niggli_cell):
        '''Calculate the linear cell variation between unit cells'''
        # Try needed to account for ref niggli cells that didn't work
        try:
            length_penalty, angle_penalty = self.penalty(ref_niggli_cell)

            total_penalty = length_penalty + angle_penalty

            return total_penalty, length_penalty, angle_penalty
        except ValueError:
            return None

    def cell_tolerance(self, ref_niggli_cell):
        '''Function to remove cells outside certain tolerances from being considered'''

        # Calculate if the 6 lattice parameters are within tolerances
        a = np.allclose(self.niggli_cell[0], ref_niggli_cell[0], atol=self.tolerance_a)
        b = np.allclose(self.niggli_cell[1], ref_niggli_cell[1], atol=self.tolerance_b)
        c = np.allclose(self.niggli_cell[2], ref_niggli_cell[2], atol=self.tolerance_c)
        alpha = np.allclose(self.niggli_cell[3], ref_niggli_cell[3], atol=self.tolerance_alpha)
        beta = np.allclose(self.niggli_cell[4], ref_niggli_cell[4], atol=self.tolerance_beta)
        gamma = np.allclose(self.niggli_cell[5], ref_niggli_cell[5], atol=self.tolerance_gamma)

        within_tol = False
        if a and b and c and alpha and beta and gamma:
            within_tol = True

        return within_tol

    def penalty(self, input_cell):
        '''Function to add penalty score depending on difference between unit cells'''
        # Split the experimental niggli cell into into component parts
        a_e, b_e, c_e, alpha_e, beta_e, gamma_e = self.niggli_cell

        # Split the input niggli cell into into component parts
        a_r, b_r, c_r, alpha_r, beta_r, gamma_r = input_cell

        delta_a = abs(a_e - a_r)
        delta_b = abs(b_e - b_r)
        delta_c = abs(c_e - c_r)
        delta_alpha = abs(alpha_e - alpha_r)
        delta_beta = abs(beta_e - beta_r)
        delta_gamma = abs(gamma_e - gamma_r)

        lengths_penalty_score = delta_a + delta_b + delta_c
        angles_penalty_score = delta_alpha + delta_beta + delta_gamma
        return lengths_penalty_score, angles_penalty_score

    def rank_scores(self):
        '''Rank the scores in self.results by Penalty score'''
        self.results = sorted(self.results, key=lambda x: float(x.Penalty_score), reverse=False)
        return

    def return_results(self):
        '''Log the results from the lattice parameter search'''
        # Output headers for the log file table
        logger.info("The lattice parameter search found the following structures:")
        logger.info(
            "PDB_CODE |   A   |   B   |   C   | ALPHA |  BETA | GAMMA | LENGTH_PENALTY | ANGLE_PENALTY | TOTAL_PENALTY ")

        count = 0
        for result in self.results:
            # Create a CSV for reading later
            with open('lattice.csv', 'a') as f:
                if count < 20:
                    f.write("{0},{1},{2},{3},{4}\n".format(result.PDB_code,
                                                           result.unit_cell,
                                                           result.Length_penalty,
                                                           result.Angle_penalty,
                                                           result.Penalty_score))
                    count += 1

            # Output a table to the log file
            uc = result.unit_cell.split(',')
            logger.info(
                '  {0}   | {1} | {2} | {3} | {4} | {5} | {6} |    {7}    |    {8}    |    {9}'.format(result.PDB_code,
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[0]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[1]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[2]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[3]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[4]))),
                                                                                ("{0:.2f}".format(
                                                                                float(
                                                                                uc[5]))),
                                                                                result.Length_penalty,
                                                                                result.Angle_penalty,
                                                                                result.Penalty_score, ))

        return

    def copy_pdbs(self, optd):
        """Copy across the top 20 PDB files identified by lattice parameter search from a local download of the PDB"""

        count = 0
        for result in self.results:
            if count <= 20:
                with gzip.open(
                        os.path.join(optd['PDB'], '{0}', 'pdb{1}.ent.gz'.format(result.PDB_code[1:3], result.PDB_code)),
                        'rb') as f:
                    file_content = f.read()
                    with open(os.path.join(optd['work_dir'], 'lattice_input_models', '{0}.pdb'.format(result.PDB_code),
                                           'w')) as o:
                        o.write(file_content)
                count += 1
        return

    def download_pdbs(self, optd):
        '''Download the top 20 results directly from the PDB'''
        pdbl = PDBList()
        count = 0
        for result in self.results:
            if count < 20:
                # Download PDB file
                pdbl.retrieve_pdb_file(result.PDB_code, pdir=os.path.join(optd['work_dir'], 'lattice_input_models'))

                # Rename the PDB file as appropriate
                os.rename(
                    os.path.join(optd['work_dir'], 'lattice_input_models', 'pdb{0}.ent'.format(result.PDB_code.lower())),
                    os.path.join(optd['work_dir'], 'lattice_input_models', '{0}.pdb'.format(result.PDB_code)))
                count += 1

        logger.info("")

        return

def return_result_list():
    result_list = []
    with open('lattice.csv', 'r') as f:
        for line in f:
            PDB_code = line.split(',')[0]
            result_list.append(PDB_code)

    return result_list<|endoftext|>


<commit_before>from django.conf import settings
from django.core.urlresolvers import reverse
from django.test.client import Client
import os

from ._utils import TestCase

from .. import models

class DonationFormViewTestCase(TestCase):
    def setUp(self):
        self.client = Client()
        settings.TEMPLATE_DIRS = (
            os.path.join(os.path.dirname(__file__), "_templates"),
        )

    @property
    def url(self):
        return reverse("donations_form")

    def get_donation_form_response(self):
        response = self.client.get(self.url)
        self.assertEqual(200, response.status_code, msg="sanity check")
        return response

    def assert_in_context(self, response, name):
        self.assertTrue(name in response.context,
                msg="%s was not in the context")

    def assert_type_in_context(self, response, name, expected_type):
        self.assert_in_context(response, name)
        self.assertEqual(response.context[name].__class__, expected_type,
                msg="%s in the context, but does not have a class of %s" % (
                        name, expected_type.__name__))

    def assert_value_in_context(self, response, name, expected_value):
        self.assert_in_context(response, name)
        self.assertEqual(response.context[name], expected_value,
                msg="%s in the context, but not equal to '%s'" % (
                        name, expected_value))

    def test_adds_form_action_url_to_context(self):
        response = self.get_donation_form_response()
        self.assert_value_in_context(response, "form_action_url", "")

    def _test_saves_donation_on_post_with_minimal_information(self):
        donor_name = self.random_donor_name
        random_amount = self.random_amount
        data = {
            "name": donor_name,
            "amount": random_amount,
        }

        # sanity check
        self.assertRaises(models.Donor.DoesNotExist,
                models.Donor.objects.get, name=donor_name)
        self.client.post(self.url, data)
        donor = models.Donor.objects.get(name=donor_name)
        self.assertEqual(donor.name, donor_name)

    def test_saves_address_if_present(self):
        donor_name = self.random_donor_name
        address_kwargs = self.random_address_kwargs
        address_formset = self.get_data_as_formset(address_kwargs)
        data = {
            "name": donor_name,
        }
        data.update(address_formset)

        self.client.post(self.url, data)
        address = models.DonorMailingAddress.objects.get(**address_kwargs)
        donor = models.Donor.objects.get(name=donor_name)
        self.assertEqual(address, donor.address)
        self.assertEqual(None, donor.mailing_address)

    def test_saves_mailing_address_if_present(self):
        donor_name = self.random_donor_name
        address_kwargs = self.random_address_kwargs
        mailing_address_kwargs = self.random_address_kwargs
        address_formset = self.get_data_as_formset([
            address_kwargs,
            mailing_address_kwargs,
        ])
        data = {
            "name": donor_name,
        }
        data.update(address_formset)

        self.assertEqual(0, len(models.DonorMailingAddress.objects.all()),
            msg="sanity check")
        self.client.post(self.url, data)
        self.assertEqual(2, len(models.DonorMailingAddress.objects.all()))
        address = models.DonorMailingAddress.objects.get(**address_kwargs)
        mailing_address = models.DonorMailingAddress.objects.get(
                **mailing_address_kwargs)
        self.assertNotEqual(address, mailing_address)

        donor = models.Donor.objects.get(name=donor_name)
        self.assertEqual(address, donor.address)
        self.assertEqual(mailing_address, donor.mailing_address)

    @property
    def random_post_data(self):
        donor_name = self.random_donor_name
        address_kwargs = self.random_address_kwargs
        address_formset = self.get_data_as_formset(address_kwargs)
        data = {
            "name": donor_name,
        }
        data.update(address_formset)
        return data

    def test_saves_mailing_address_if_same_as_billing_is_checked(self):
        data = self.random_post_data
        data["mailing_same_as_billing"] = u"1"
        self.client.post(self.url, data)
        donor = models.Donor.objects.get(name=data["name"])
        self.assertEqual(donor.address, donor.mailing_address)
<commit_msg>add in a couple of test cases that need to be written<commit_after>from django.conf import settings
from django.core.urlresolvers import reverse
from django.test.client import Client
import os
from unittest import expectedFailure

from ._utils import TestCase

from .. import models

class DonationFormViewTestCase(TestCase):
    def setUp(self):
        self.client = Client()
        settings.TEMPLATE_DIRS = (
            os.path.join(os.path.dirname(__file__), "_templates"),
        )

    @property
    def url(self):
        return reverse("donations_form")

    def get_donation_form_response(self):
        response = self.client.get(self.url)
        self.assertEqual(200, response.status_code, msg="sanity check")
        return response

    def assert_in_context(self, response, name):
        self.assertTrue(name in response.context,
                msg="%s was not in the context")

    def assert_type_in_context(self, response, name, expected_type):
        self.assert_in_context(response, name)
        self.assertEqual(response.context[name].__class__, expected_type,
                msg="%s in the context, but does not have a class of %s" % (
                        name, expected_type.__name__))

    def assert_value_in_context(self, response, name, expected_value):
        self.assert_in_context(response, name)
        self.assertEqual(response.context[name], expected_value,
                msg="%s in the context, but not equal to '%s'" % (
                        name, expected_value))

    def test_adds_form_action_url_to_context(self):
        response = self.get_donation_form_response()
        self.assert_value_in_context(response, "form_action_url", "")

    def _test_saves_donation_on_post_with_minimal_information(self):
        donor_name = self.random_donor_name
        random_amount = self.random_amount
        data = {
            "name": donor_name,
            "amount": random_amount,
        }

        # sanity check
        self.assertRaises(models.Donor.DoesNotExist,
                models.Donor.objects.get, name=donor_name)
        self.client.post(self.url, data)
        donor = models.Donor.objects.get(name=donor_name)
        self.assertEqual(donor.name, donor_name)

    def test_saves_address_if_present(self):
        donor_name = self.random_donor_name
        address_kwargs = self.random_address_kwargs
        address_formset = self.get_data_as_formset(address_kwargs)
        data = {
            "name": donor_name,
        }
        data.update(address_formset)

        self.client.post(self.url, data)
        address = models.DonorMailingAddress.objects.get(**address_kwargs)
        donor = models.Donor.objects.get(name=donor_name)
        self.assertEqual(address, donor.address)
        self.assertEqual(None, donor.mailing_address)

    def test_saves_mailing_address_if_present(self):
        donor_name = self.random_donor_name
        address_kwargs = self.random_address_kwargs
        mailing_address_kwargs = self.random_address_kwargs
        address_formset = self.get_data_as_formset([
            address_kwargs,
            mailing_address_kwargs,
        ])
        data = {
            "name": donor_name,
        }
        data.update(address_formset)

        self.assertEqual(0, len(models.DonorMailingAddress.objects.all()),
            msg="sanity check")
        self.client.post(self.url, data)
        self.assertEqual(2, len(models.DonorMailingAddress.objects.all()))
        address = models.DonorMailingAddress.objects.get(**address_kwargs)
        mailing_address = models.DonorMailingAddress.objects.get(
                **mailing_address_kwargs)
        self.assertNotEqual(address, mailing_address)

        donor = models.Donor.objects.get(name=donor_name)
        self.assertEqual(address, donor.address)
        self.assertEqual(mailing_address, donor.mailing_address)

    @property
    def random_post_data(self):
        donor_name = self.random_donor_name
        address_kwargs = self.random_address_kwargs
        address_formset = self.get_data_as_formset(address_kwargs)
        data = {
            "name": donor_name,
        }
        data.update(address_formset)
        return data

    def test_saves_mailing_address_if_same_as_billing_is_checked(self):
        data = self.random_post_data
        data["mailing_same_as_billing"] = u"1"
        self.client.post(self.url, data)
        donor = models.Donor.objects.get(name=data["name"])
        self.assertEqual(donor.address, donor.mailing_address)

    @expectedFailure
    def test_saves_donation_information(self):
        self.fail()

    @expectedFailure
    def test_displays_errors_on_donor_validation_error(self):
        self.fail()

    @expectedFailure
    def test_displays_errors_on_address_validation_error(self):
        self.fail()

    @expectedFailure
    def test_redirects_to_success_url_on_success(self):
        self.fail()

    @expectedFailure
    def test_displays_errors_when_payment_method_authorization_fails(self):
        self.fail()
<|endoftext|>


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

# suggest commonly co-used arguments
# search through shell history
# screenshots in case explain shell goes down, fork explain shell
import argparse
import requests
from requests.exceptions import ConnectionError
from urllib import quote
from pyquery import PyQuery
from clint.textui import puts, indent, colored

def request_page(url):
    #print url
    # set User-Agent in header to receive appropriate response
    user_agent = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.7; rv:11.0) Gecko/20100101 Firefox/11.0'
    return requests.get(url, headers={'User-Agent': user_agent}).text

def retrieve_so_page_links(query):
    """Search within stackoverflow (SO) using google and retrieve links"""
    # augment query terms
    query += " linux command line"
    search_url = 'http://www.google.com/search?q=site:stackoverflow.com%20{0}'
    url = search_url.format(quote(query))
    result = request_page(url)
    html = PyQuery(result)
    return [a.attrib['href'] for a in html('.l')] or \
           [a.attrib['href'] for a in html('.r')('a')]

def get_command_line(args, link):
    if not link:
        return False
    page = request_page(link + '?answertab=votes')
    html = PyQuery(page)

    commands = [l.strip() for l in open('nix_commands.txt').readlines()]
    cmd_frequency = {}
    cmd_lines = {}
    max_cmd = None
    max_freq = 0
    html_answers = html('.answercell').find('*')
    for line in html_answers:
        if line is not None and line.text is not None and not line.text.isspace() and len(line.text.strip().split()) > 1:
            line_text = line.text.strip('$>').strip().split('\n')[0]

            cmd = line.text.strip('$>').strip().split()[0]
            if cmd in commands:
                cmd_frequency[cmd] = cmd_frequency.get(cmd, 0) + 1
                cmd_lines.setdefault(cmd, []).append(line_text)
                if cmd_frequency > max_freq:
                    max_freq = cmd_frequency
                    max_cmd = cmd
    if max_cmd is not None:
        return cmd_lines[max_cmd][0]

    # TODO handle case in which None is returned
    return None

def get_explanation_lines(command_line):
    if command_line is None:
        return None
    explain_url = 'http://www.explainshell.com/explain?cmd={0}'
    url = explain_url.format(quote(command_line))
    result = request_page(url)
    html = PyQuery(result)
    return [PyQuery(hb).text() for hb in html('.help-box')]

def print_result(command_line, explanation_lines):
    if command_line and len(explanation_lines) > 1:
        title, body = explanation_lines[0], explanation_lines[1:]
        with indent(4, quote=colored.yellow('|')):
            puts()
            puts(colored.yellow(title.upper()))
            with indent(4):
                for line in body:
                    puts(line.encode('ascii', 'ignore'))
            puts()
            with indent(2, quote=colored.green('>')):
                puts(colored.green(command_line.encode('ascii', 'ignore')))
            puts()
    else:
        with indent(4, quote=colored.magenta('|')):
            puts()
            puts(colored.magenta("sorry, I don't understand your query, try a different query"))
            puts()

        # TODO remove DEBUG print statements
        #print command_line
        #print explanation_lines

def hownix(args):
    try:
        query = ' '.join(args['query']).replace('?', '')
        links = retrieve_so_page_links(query)
        if not links:
            return False
        command_line = get_command_line(args, links[0])
        explanation_lines = get_explanation_lines(command_line)
        print_result(command_line, explanation_lines)
    except ConnectionError:
        print 'Unable to connect to service'

def main():
    parser = argparse.ArgumentParser(description='find and understand *nix commands')
    parser.add_argument('query', metavar='QUERY', type=str, nargs='*',
                        help='query')
    args = vars(parser.parse_args())
    if not args['query']:
        parser.print_help()
        return
    hownix(args)

if __name__ == '__main__':
    main()
<commit_msg>cleanup, add comments<commit_after>#!/usr/bin/env python

# TODO:
# - suggest commonly co-used arguments
# - search through shell history

import argparse
import requests
from requests.exceptions import ConnectionError
from urllib import quote
from pyquery import PyQuery
from clint.textui import puts, indent, colored

def request_page(url):
    """Returns a http response object corresponding to url"""
    # set User-Agent in header to receive appropriate response
    user_agent = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.7; rv:11.0) Gecko/20100101 Firefox/11.0'
    return requests.get(url, headers={'User-Agent': user_agent}).text

def retrieve_so_page_links(query):
    """Search within stackoverflow (SO) using google and retrieve links"""
    # augment query terms
    query += " linux command line"
    search_url = 'http://www.google.com/search?q=site:stackoverflow.com%20{0}'
    url = search_url.format(quote(query))
    result = request_page(url)
    html = PyQuery(result)
    return [a.attrib['href'] for a in html('.l')] or \
           [a.attrib['href'] for a in html('.r')('a')]

def get_command_line(args, links):
    """Parse through StackOverflow page and identify command line corresponding to most frequently referred to command"""
    if not links:
        return False
    link = links[0] # use the first link
    page = request_page(link + '?answertab=votes')
    html = PyQuery(page)
    commands = [l.strip() for l in open('nix_commands.txt').readlines()]
    cmd_frequency = {}
    cmd_lines = {}
    max_cmd = None
    max_freq = 0
    html_answers = html('.answercell').find('*')
    for line in html_answers:
        if line.text is not None and not line.text.isspace() and len(line.text.strip().split()) > 1:
            line_text = line.text.strip('$>').strip().split('\n')[0]
            cmd = line.text.split()[0]
            if cmd in commands:
                cmd_frequency[cmd] = cmd_frequency.get(cmd, 0) + 1
                cmd_lines.setdefault(cmd, []).append(line_text)
                if cmd_frequency > max_freq:
                    max_freq = cmd_frequency
                    max_cmd = cmd
    if max_cmd is not None:
        return cmd_lines[max_cmd][0]
    return None

def get_explanation_lines(command_line):
    """Returns list of text lines explaining command line"""
    if command_line is None:
        return None
    explain_url = 'http://www.explainshell.com/explain?cmd={0}'
    url = explain_url.format(quote(command_line))
    result = request_page(url)
    html = PyQuery(result)
    return [PyQuery(hb).text() for hb in html('.help-box')]

def print_result(command_line, explanation_lines):
    """Prints sample command line preceded by explanation"""
    if command_line and len(explanation_lines) > 1:
        title, body = explanation_lines[0], explanation_lines[1:]
        with indent(4, quote=colored.yellow('|')):
            puts()
            puts(colored.yellow(title.upper()))
            with indent(4):
                for line in body:
                    puts(line.encode('ascii', 'ignore'))
            puts()
            with indent(2, quote=colored.green('>')):
                puts(colored.green(command_line.encode('ascii', 'ignore')))
            puts()
    else:
        with indent(4, quote=colored.magenta('|')):
            puts()
            puts(colored.magenta("couldn't find an answer, try a different query"))
            puts()

def hownix(args):
    """Find StackOverflow pages related to query via Google. Extract command lines and explanation"""
    try:
        query = ' '.join(args['query']).replace('?', '')
        links = retrieve_so_page_links(query)
        command_line = get_command_line(args, links)
        explanation_lines = get_explanation_lines(command_line)
        print_result(command_line, explanation_lines)
    except ConnectionError:
        puts(colored.red('Unable to connect to service'))

def main():
    parser = argparse.ArgumentParser(description='find and understand *nix commands')
    parser.add_argument('query', nargs='*', help='query')
    args = vars(parser.parse_args())
    if not args['query']:
        parser.print_help()
        return
    hownix(args)

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