<commit_before>import numpy as np

from opensfm import context
from opensfm import types
from opensfm import pygeometry

"""
Trying to imitate the following structure

reconstruction = {
    "cameras": {
        "theta": {
            "projection_type": "equirectangular"
        }
    },

    "shots" : {
        'im1': {
            "camera": "theta",
            "rotation": [0.0, 0.0, 0.0],
            "translation": [0.0, 0.0, 0.0],
        },
        'im2': {
            "camera": "theta",
            "rotation": [0, 0, 0.0],
            "translation": [-1, 0, 0.0],
        },
    },

    "points" : {
    },
}
"""

def test_reconstruction_class_initialization():

    # Instantiate Reconstruction
    reconstruction = types.Reconstruction()

    # Instantiate camera instrinsics
    camera = types.PerspectiveCamera()
    camera.id = 'apple iphone 4s back camera 4.28mm f/2.4'
    camera.focal = 0.9722222222222222
    camera.k1 = 0.006094395128698237
    camera.k2 = -0.0004952058188617129
    camera.height = 2448
    camera.width = 3264

    # Instantiate GPS data
    metadata = types.ShotMetadata()
    metadata.orientation = 1
    metadata.capture_time = 0.0
    metadata.gps_dop = 5.0
    metadata.gps_position = [1.0815875281451939,
                             -0.96510451436708888,
                             1.2042133903991235]

    # Instantiate shots
    pose0 = types.Pose([0.0, 0.0, 0.0], [0.0, 0.0, 0.0])

    shot0 = types.Shot()
    shot0.id = 0
    shot0.camera = camera
    shot0.pose = pose0
    shot0.metadata = metadata

    pose1 = types.Pose([0.0, 0.0, 0.0], [-1.0, 0.0, 0.0])

    shot1 = types.Shot()
    shot1.id = 1
    shot1.camera = camera
    shot1.pose = pose1
    shot1.metadata = metadata

    # Add info to current reconstruction
    reconstruction.add_camera(camera)
    reconstruction.add_shot(shot0)
    reconstruction.add_shot(shot1)

    # TEST
    assert len(reconstruction.cameras) == 1
    assert len(reconstruction.shots) == 2
    assert len(reconstruction.points) == 0
    assert reconstruction.get_camera(camera.id) == camera
    assert reconstruction.get_camera(1) is None
    assert reconstruction.get_shot(shot0.id) == shot0
    assert reconstruction.get_shot(shot1.id) == shot1
    assert reconstruction.get_shot(2) is None

def test_perspective_camera_projection():
    """Test perspectiive projection--backprojection loop."""
    for camera in _get_perspective_camera():
        pixel = [0.1, 0.2]
        bearing = camera.pixel_bearing(pixel)
        projected = camera.project(bearing)
        assert np.allclose(pixel, projected)

def test_fisheye_camera_projection():
    """Test fisheye projection--backprojection loop."""
    if not context.OPENCV3:
        return
    for camera in _get_fisheye_camera():
        pixel = [0.1, 0.2]
        bearing = camera.pixel_bearing(pixel)
        projected = camera.project(bearing)
        assert np.allclose(pixel, projected)

def test_dual_camera_projection():
    """Test dual projection--backprojection loop."""
    if not context.OPENCV3:
        return
    for camera in _get_dual_camera():
        pixel = [0.1, 0.2]
        bearing = camera.pixel_bearing(pixel)
        projected = camera.project(bearing)
        assert np.allclose(pixel, projected)

def test_spherical_camera_projection():
    """Test spherical projection--backprojection loop."""
    for camera in _get_spherical_camera():
        pixel = [0.1, 0.2]
        bearing = camera.pixel_bearing(pixel)
        projected = camera.project(bearing)
        assert np.allclose(pixel, projected)

def test_pose_properties():
    """Test pose constructor, getters and setters."""
    p = types.Pose([1, 2, 3], [4, 5, 6])
    assert np.allclose(p.rotation, [1, 2, 3])
    assert type(p.rotation) == np.ndarray
    assert p.rotation.dtype == float
    assert np.allclose(p.translation, [4, 5, 6])
    assert type(p.translation) == np.ndarray
    assert p.translation.dtype == float

def test_pose_inverse():
    p = types.Pose([1, 2, 3], [4, 5, 6])
    inverse = p.inverse()
    identity = p.compose(inverse)
    assert np.allclose(identity.rotation, [0, 0, 0])
    assert np.allclose(identity.translation, [0, 0, 0])

def test_single_vs_many():
    points = np.array([[1, 2, 3], [4, 5, 6]], dtype=float)
    pixels = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=float)
    depths = np.array([1, 2], dtype=float)

    pose = types.Pose([1, 2, 3], [4, 5, 6])
    t_single = [pose.transform(p) for p in points]
    t_many = pose.transform_many(points)
    assert np.allclose(t_single, t_many)

    t_single = [pose.transform_inverse(p) for p in points]
    t_many = pose.transform_inverse_many(points)
    assert np.allclose(t_single, t_many)

    cameras = [
        _get_perspective_camera(),
        _get_brown_perspective_camera(),
        _get_spherical_camera(),
    ]
    if context.OPENCV3:
        cameras.append(_get_fisheye_camera())

    for camera, camera_cpp in cameras:
        p = camera.project_many(points)
        p_cpp = camera_cpp.project_many(points)
        assert np.allclose(p, p_cpp)

        b = camera.pixel_bearing_many(pixels)
        b_cpp = camera_cpp.pixel_bearing_many(pixels)
        print(camera)
        assert np.allclose(b, b_cpp)

        if hasattr(camera, 'back_project'):
            q_single = [camera.back_project(p, d)
                        for p, d in zip(pixels, depths)]
            q_many = camera.back_project_many(pixels, depths)
            assert np.allclose(q_single, q_many)

def _get_perspective_camera():
    camera = types.PerspectiveCamera()
    camera.width = 800
    camera.height = 600
    camera.focal = 0.6
    camera.k1 = -0.1
    camera.k2 = 0.01
    camera_cpp = pygeometry.Camera.create_perspective(
        camera.focal, camera.k1, camera.k2)
    return camera, camera_cpp

def _get_brown_perspective_camera():
    camera = types.BrownPerspectiveCamera()
    camera.width = 800
    camera.height = 600
    camera.focal_x = 0.6
    camera.focal_y = 0.7
    camera.c_x = 0.1
    camera.c_y = -0.05
    camera.k1 = -0.1
    camera.k2 = 0.01
    camera.p1 = 0.001
    camera.p2 = 0.002
    camera.k3 = 0.01
    camera_cpp = pygeometry.Camera.create_brown(
        camera.focal_x, camera.focal_y / camera.focal_x,
        [camera.c_x, camera.c_y],
        [camera.k1, camera.k2, camera.k3, camera.p1, camera.p2])
    return camera, camera_cpp

def _get_fisheye_camera():
    camera = types.FisheyeCamera()
    camera.width = 800
    camera.height = 600
    camera.focal = 0.6
    camera.k1 = -0.1
    camera.k2 = 0.01
    camera_cpp = pygeometry.Camera.create_fisheye(
        camera.focal, camera.k1, camera.k2)
    return camera, camera_cpp

def _get_dual_camera():
    camera = types.DualCamera()
    camera.width = 800
    camera.height = 600
    camera.focal = 0.3
    camera.k1 = -0.1
    camera.k2 = 0.01
    camera.transition = 0.5
    camera_cpp = pygeometry.Camera.create_dual(
        camera.transition, camera.focal, camera.k1, camera.k2)
    return camera, camera_cpp

def _get_spherical_camera():
    camera = types.SphericalCamera()
    camera.width = 800
    camera.height = 600
    camera_cpp = pygeometry.Camera.create_spherical()
    return camera, camera_cpp
<commit_msg>fix: removed print<commit_after>import numpy as np

from opensfm import context
from opensfm import types
from opensfm import pygeometry

"""
Trying to imitate the following structure

reconstruction = {
    "cameras": {
        "theta": {
            "projection_type": "equirectangular"
        }
    },

    "shots" : {
        'im1': {
            "camera": "theta",
            "rotation": [0.0, 0.0, 0.0],
            "translation": [0.0, 0.0, 0.0],
        },
        'im2': {
            "camera": "theta",
            "rotation": [0, 0, 0.0],
            "translation": [-1, 0, 0.0],
        },
    },

    "points" : {
    },
}
"""

def test_reconstruction_class_initialization():

    # Instantiate Reconstruction
    reconstruction = types.Reconstruction()

    # Instantiate camera instrinsics
    camera = types.PerspectiveCamera()
    camera.id = 'apple iphone 4s back camera 4.28mm f/2.4'
    camera.focal = 0.9722222222222222
    camera.k1 = 0.006094395128698237
    camera.k2 = -0.0004952058188617129
    camera.height = 2448
    camera.width = 3264

    # Instantiate GPS data
    metadata = types.ShotMetadata()
    metadata.orientation = 1
    metadata.capture_time = 0.0
    metadata.gps_dop = 5.0
    metadata.gps_position = [1.0815875281451939,
                             -0.96510451436708888,
                             1.2042133903991235]

    # Instantiate shots
    pose0 = types.Pose([0.0, 0.0, 0.0], [0.0, 0.0, 0.0])

    shot0 = types.Shot()
    shot0.id = 0
    shot0.camera = camera
    shot0.pose = pose0
    shot0.metadata = metadata

    pose1 = types.Pose([0.0, 0.0, 0.0], [-1.0, 0.0, 0.0])

    shot1 = types.Shot()
    shot1.id = 1
    shot1.camera = camera
    shot1.pose = pose1
    shot1.metadata = metadata

    # Add info to current reconstruction
    reconstruction.add_camera(camera)
    reconstruction.add_shot(shot0)
    reconstruction.add_shot(shot1)

    # TEST
    assert len(reconstruction.cameras) == 1
    assert len(reconstruction.shots) == 2
    assert len(reconstruction.points) == 0
    assert reconstruction.get_camera(camera.id) == camera
    assert reconstruction.get_camera(1) is None
    assert reconstruction.get_shot(shot0.id) == shot0
    assert reconstruction.get_shot(shot1.id) == shot1
    assert reconstruction.get_shot(2) is None

def test_perspective_camera_projection():
    """Test perspectiive projection--backprojection loop."""
    for camera in _get_perspective_camera():
        pixel = [0.1, 0.2]
        bearing = camera.pixel_bearing(pixel)
        projected = camera.project(bearing)
        assert np.allclose(pixel, projected)

def test_fisheye_camera_projection():
    """Test fisheye projection--backprojection loop."""
    if not context.OPENCV3:
        return
    for camera in _get_fisheye_camera():
        pixel = [0.1, 0.2]
        bearing = camera.pixel_bearing(pixel)
        projected = camera.project(bearing)
        assert np.allclose(pixel, projected)

def test_dual_camera_projection():
    """Test dual projection--backprojection loop."""
    if not context.OPENCV3:
        return
    for camera in _get_dual_camera():
        pixel = [0.1, 0.2]
        bearing = camera.pixel_bearing(pixel)
        projected = camera.project(bearing)
        assert np.allclose(pixel, projected)

def test_spherical_camera_projection():
    """Test spherical projection--backprojection loop."""
    for camera in _get_spherical_camera():
        pixel = [0.1, 0.2]
        bearing = camera.pixel_bearing(pixel)
        projected = camera.project(bearing)
        assert np.allclose(pixel, projected)

def test_pose_properties():
    """Test pose constructor, getters and setters."""
    p = types.Pose([1, 2, 3], [4, 5, 6])
    assert np.allclose(p.rotation, [1, 2, 3])
    assert type(p.rotation) == np.ndarray
    assert p.rotation.dtype == float
    assert np.allclose(p.translation, [4, 5, 6])
    assert type(p.translation) == np.ndarray
    assert p.translation.dtype == float

def test_pose_inverse():
    p = types.Pose([1, 2, 3], [4, 5, 6])
    inverse = p.inverse()
    identity = p.compose(inverse)
    assert np.allclose(identity.rotation, [0, 0, 0])
    assert np.allclose(identity.translation, [0, 0, 0])

def test_single_vs_many():
    points = np.array([[1, 2, 3], [4, 5, 6]], dtype=float)
    pixels = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=float)
    depths = np.array([1, 2], dtype=float)

    pose = types.Pose([1, 2, 3], [4, 5, 6])
    t_single = [pose.transform(p) for p in points]
    t_many = pose.transform_many(points)
    assert np.allclose(t_single, t_many)

    t_single = [pose.transform_inverse(p) for p in points]
    t_many = pose.transform_inverse_many(points)
    assert np.allclose(t_single, t_many)

    cameras = [
        _get_perspective_camera(),
        _get_brown_perspective_camera(),
        _get_spherical_camera(),
    ]
    if context.OPENCV3:
        cameras.append(_get_fisheye_camera())

    for camera, camera_cpp in cameras:
        p = camera.project_many(points)
        p_cpp = camera_cpp.project_many(points)
        assert np.allclose(p, p_cpp)

        b = camera.pixel_bearing_many(pixels)
        b_cpp = camera_cpp.pixel_bearing_many(pixels)
        assert np.allclose(b, b_cpp)

        if hasattr(camera, 'back_project'):
            q_single = [camera.back_project(p, d)
                        for p, d in zip(pixels, depths)]
            q_many = camera.back_project_many(pixels, depths)
            assert np.allclose(q_single, q_many)

def _get_perspective_camera():
    camera = types.PerspectiveCamera()
    camera.width = 800
    camera.height = 600
    camera.focal = 0.6
    camera.k1 = -0.1
    camera.k2 = 0.01
    camera_cpp = pygeometry.Camera.create_perspective(
        camera.focal, camera.k1, camera.k2)
    return camera, camera_cpp

def _get_brown_perspective_camera():
    camera = types.BrownPerspectiveCamera()
    camera.width = 800
    camera.height = 600
    camera.focal_x = 0.6
    camera.focal_y = 0.7
    camera.c_x = 0.1
    camera.c_y = -0.05
    camera.k1 = -0.1
    camera.k2 = 0.01
    camera.p1 = 0.001
    camera.p2 = 0.002
    camera.k3 = 0.01
    camera_cpp = pygeometry.Camera.create_brown(
        camera.focal_x, camera.focal_y / camera.focal_x,
        [camera.c_x, camera.c_y],
        [camera.k1, camera.k2, camera.k3, camera.p1, camera.p2])
    return camera, camera_cpp

def _get_fisheye_camera():
    camera = types.FisheyeCamera()
    camera.width = 800
    camera.height = 600
    camera.focal = 0.6
    camera.k1 = -0.1
    camera.k2 = 0.01
    camera_cpp = pygeometry.Camera.create_fisheye(
        camera.focal, camera.k1, camera.k2)
    return camera, camera_cpp

def _get_dual_camera():
    camera = types.DualCamera()
    camera.width = 800
    camera.height = 600
    camera.focal = 0.3
    camera.k1 = -0.1
    camera.k2 = 0.01
    camera.transition = 0.5
    camera_cpp = pygeometry.Camera.create_dual(
        camera.transition, camera.focal, camera.k1, camera.k2)
    return camera, camera_cpp

def _get_spherical_camera():
    camera = types.SphericalCamera()
    camera.width = 800
    camera.height = 600
    camera_cpp = pygeometry.Camera.create_spherical()
    return camera, camera_cpp
<|endoftext|>


<commit_before>import numpy as np
import mdtraj as md
from unittest import skipIf
import logging
from mdtraj.testing import eq
from openmoltools import utils
import simtk.unit as u
from simtk.openmm import app
import simtk.openmm as mm
import simtk.openmm.openmm as mmmm
from distutils.spawn import find_executable
import parmed

HAVE_RDKIT = True
try:
    from rdkit import Chem
    from rdkit.Chem import AllChem
except ImportError:
    HAVE_RDKIT = False

OBABEL_PATH = find_executable("obabel")
SKIP_SMILES = not ((HAVE_RDKIT) & (OBABEL_PATH is not None))

CHECKMOL_PATH = find_executable("checkmol")
SKIP_CHECKMOL = (CHECKMOL_PATH is None)

logging.basicConfig(level=logging.DEBUG, format="LOG: %(message)s")

def test_enter_temp_directory():
    with utils.enter_temp_directory():
        pass

def test_parse_ligand_filename():
    molecule_name = "sustiva"
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    name, ext = utils.parse_ligand_filename(input_filename)

    eq(name, "sustiva")
    eq(ext, ".mol2")

def test_run_test_molecule():
    molecule_name = "sustiva"
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    with utils.enter_temp_directory():  # Prevents creating tons of GAFF files everywhere.
        utils.test_molecule(molecule_name, input_filename)

def test_acpype_conversion():
    molecule_name = 'sustiva'
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    with utils.enter_temp_directory(): # Prevents creating tons of GAFF files everywhere.
        gaff_mol2_filename, frcmod_filename = amber.run_antechamber(molecule_name, input_filename, charge_method=None)
        prmtop, inpcrd = utils.run_tleap(molecule_name, gaff_mol2_filename, frcmod_filename)
        out_top, out_gro = utils.convert_via_acpype( molecule_name, prmtop, inpcrd )

def test_parmed_conversion():
    molecule_name = 'sustiva'
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    with utils.enter_temp_directory(): # Prevents creating tons of GAFF files everywhere.
        #Make sure conversion runs
        gaff_mol2_filename, frcmod_filename = amber.run_antechamber(molecule_name, input_filename, charge_method=None)
        prmtop, inpcrd = utils.run_tleap(molecule_name, gaff_mol2_filename, frcmod_filename)
        out_top, out_gro = utils.amber_to_gromacs( molecule_name, prmtop, inpcrd, precision = 8 )

        #Test energies before and after conversion
        #Set up amber system
        a = parmed.amber.AmberParm( prmtop, inpcrd )
        ambersys = a.createSystem()
        ambercon = mmmm.Context( ambersys, mm.VerletIntegrator(0.001))
        ambercon.setPositions( a.positions )
        #Set up GROMACS system
        g = parmed.load_file( out_top )
        gro = parmed.gromacs.GromacsGroFile.parse( out_gro )
        g.box = gro.box
        g.positions = gro.positions
        gromacssys = g.createSystem()
        gromacscon = mmmm.Context( gromacssys, mm.VerletIntegrator(0.001))
        gromacscon.setPositions( g.positions )

        #Check energies
        a_energies = parmed.openmm.utils.energy_decomposition( a, ambercon )
        g_energies = parmed.openmm.utils.energy_decomposition( g, gromacscon )
        #Check components
        tolerance = 1e-5
        ok = True
        for key in a_energies.keys():
            diff = np.abs(a_energies[key] - g_energies[key] )
            if diff/np.abs(a_energies[key]) > tolerance:
                ok = False
                print("In testing AMBER to GROMACS conversion, %s energy differs by %.5g, which is more than a fraction %.2g of the total, so conversion appears not to be working properly." % ( key, diff, tolerance) )
        if not ok:
            raise(ValueError("AMBER to GROMACS conversion yields energies which are too different."))

@skipIf(SKIP_CHECKMOL, "Skipping testing of checkmol descriptors since checkmol is not found (under that name)." )
def test_checkmol_descriptors():
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    utils.get_checkmol_descriptors( input_filename )

@skipIf(SKIP_SMILES, "Skipping testing of smiles conversion because openbabel or rdkit not found.")
def test_smiles_conversion():
    pdb_filename = utils.get_data_filename("chemicals/proteins/1vii.pdb")
    smiles = 'Cc1ccccc1'  # Also known as toluene.

    temperature = 300 * u.kelvin
    friction = 0.3 / u.picosecond
    timestep = 0.01 * u.femtosecond

    protein_traj = md.load(pdb_filename)
    protein_traj.center_coordinates()

    protein_top = protein_traj.top.to_openmm()
    protein_xyz = protein_traj.openmm_positions(0)

    ligand_trajectories, ffxml = utils.smiles_to_mdtraj_ffxml([smiles])
    ligand_traj = ligand_trajectories[0]
    ligand_traj.center_coordinates()

    eq(ligand_traj.n_atoms, 15)
    eq(ligand_traj.n_frames, 1)

    #Move the pre-centered ligand sufficiently far away from the protein to avoid a clash.
    min_atom_pair_distance = ((ligand_traj.xyz[0] ** 2.).sum(1) ** 0.5).max() + ((protein_traj.xyz[0] ** 2.).sum(1) ** 0.5).max() + 0.3
    ligand_traj.xyz += np.array([1.0, 0.0, 0.0]) * min_atom_pair_distance

    ligand_xyz = ligand_traj.openmm_positions(0)
    ligand_top = ligand_traj.top.to_openmm()

    forcefield = app.ForceField("amber10.xml", ffxml, "tip3p.xml")

    model = app.modeller.Modeller(protein_top, protein_xyz)
    model.add(ligand_top, ligand_xyz)
    model.addSolvent(forcefield, padding=0.4 * u.nanometer)

    system = forcefield.createSystem(model.topology, nonbondedMethod=app.PME, nonbondedCutoff=1.0 * u.nanometers, constraints=app.HAngles)

    integrator = mm.LangevinIntegrator(temperature, friction, timestep)

    simulation = app.Simulation(model.topology, system, integrator)
    simulation.context.setPositions(model.positions)
    print("running")
    simulation.step(1)
<commit_msg>Fixing an import problem in test_utils.py<commit_after>import numpy as np
import mdtraj as md
from unittest import skipIf
import logging
from mdtraj.testing import eq
from openmoltools import utils
from openmoltools import amber
import simtk.unit as u
from simtk.openmm import app
import simtk.openmm as mm
import simtk.openmm.openmm as mmmm
from distutils.spawn import find_executable
import parmed

HAVE_RDKIT = True
try:
    from rdkit import Chem
    from rdkit.Chem import AllChem
except ImportError:
    HAVE_RDKIT = False

OBABEL_PATH = find_executable("obabel")
SKIP_SMILES = not ((HAVE_RDKIT) & (OBABEL_PATH is not None))

CHECKMOL_PATH = find_executable("checkmol")
SKIP_CHECKMOL = (CHECKMOL_PATH is None)

logging.basicConfig(level=logging.DEBUG, format="LOG: %(message)s")

def test_enter_temp_directory():
    with utils.enter_temp_directory():
        pass

def test_parse_ligand_filename():
    molecule_name = "sustiva"
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    name, ext = utils.parse_ligand_filename(input_filename)

    eq(name, "sustiva")
    eq(ext, ".mol2")

def test_run_test_molecule():
    molecule_name = "sustiva"
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    with utils.enter_temp_directory():  # Prevents creating tons of GAFF files everywhere.
        utils.test_molecule(molecule_name, input_filename)

def test_acpype_conversion():
    molecule_name = 'sustiva'
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    with utils.enter_temp_directory(): # Prevents creating tons of GAFF files everywhere.
        gaff_mol2_filename, frcmod_filename = amber.run_antechamber(molecule_name, input_filename, charge_method=None)
        prmtop, inpcrd = amber.run_tleap(molecule_name, gaff_mol2_filename, frcmod_filename)
        out_top, out_gro = utils.convert_via_acpype( molecule_name, prmtop, inpcrd )

def test_parmed_conversion():
    molecule_name = 'sustiva'
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    with utils.enter_temp_directory(): # Prevents creating tons of GAFF files everywhere.
        #Make sure conversion runs
        gaff_mol2_filename, frcmod_filename = amber.run_antechamber(molecule_name, input_filename, charge_method=None)
        prmtop, inpcrd = amber.run_tleap(molecule_name, gaff_mol2_filename, frcmod_filename)
        out_top, out_gro = utils.amber_to_gromacs( molecule_name, prmtop, inpcrd, precision = 8 )

        #Test energies before and after conversion
        #Set up amber system
        a = parmed.amber.AmberParm( prmtop, inpcrd )
        ambersys = a.createSystem()
        ambercon = mmmm.Context( ambersys, mm.VerletIntegrator(0.001))
        ambercon.setPositions( a.positions )
        #Set up GROMACS system
        g = parmed.load_file( out_top )
        gro = parmed.gromacs.GromacsGroFile.parse( out_gro )
        g.box = gro.box
        g.positions = gro.positions
        gromacssys = g.createSystem()
        gromacscon = mmmm.Context( gromacssys, mm.VerletIntegrator(0.001))
        gromacscon.setPositions( g.positions )

        #Check energies
        a_energies = parmed.openmm.utils.energy_decomposition( a, ambercon )
        g_energies = parmed.openmm.utils.energy_decomposition( g, gromacscon )
        #Check components
        tolerance = 1e-5
        ok = True
        for key in a_energies.keys():
            diff = np.abs(a_energies[key] - g_energies[key] )
            if diff/np.abs(a_energies[key]) > tolerance:
                ok = False
                print("In testing AMBER to GROMACS conversion, %s energy differs by %.5g, which is more than a fraction %.2g of the total, so conversion appears not to be working properly." % ( key, diff, tolerance) )
        if not ok:
            raise(ValueError("AMBER to GROMACS conversion yields energies which are too different."))

@skipIf(SKIP_CHECKMOL, "Skipping testing of checkmol descriptors since checkmol is not found (under that name)." )
def test_checkmol_descriptors():
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    utils.get_checkmol_descriptors( input_filename )

@skipIf(SKIP_SMILES, "Skipping testing of smiles conversion because openbabel or rdkit not found.")
def test_smiles_conversion():
    pdb_filename = utils.get_data_filename("chemicals/proteins/1vii.pdb")
    smiles = 'Cc1ccccc1'  # Also known as toluene.

    temperature = 300 * u.kelvin
    friction = 0.3 / u.picosecond
    timestep = 0.01 * u.femtosecond

    protein_traj = md.load(pdb_filename)
    protein_traj.center_coordinates()

    protein_top = protein_traj.top.to_openmm()
    protein_xyz = protein_traj.openmm_positions(0)

    ligand_trajectories, ffxml = utils.smiles_to_mdtraj_ffxml([smiles])
    ligand_traj = ligand_trajectories[0]
    ligand_traj.center_coordinates()

    eq(ligand_traj.n_atoms, 15)
    eq(ligand_traj.n_frames, 1)

    #Move the pre-centered ligand sufficiently far away from the protein to avoid a clash.
    min_atom_pair_distance = ((ligand_traj.xyz[0] ** 2.).sum(1) ** 0.5).max() + ((protein_traj.xyz[0] ** 2.).sum(1) ** 0.5).max() + 0.3
    ligand_traj.xyz += np.array([1.0, 0.0, 0.0]) * min_atom_pair_distance

    ligand_xyz = ligand_traj.openmm_positions(0)
    ligand_top = ligand_traj.top.to_openmm()

    forcefield = app.ForceField("amber10.xml", ffxml, "tip3p.xml")

    model = app.modeller.Modeller(protein_top, protein_xyz)
    model.add(ligand_top, ligand_xyz)
    model.addSolvent(forcefield, padding=0.4 * u.nanometer)

    system = forcefield.createSystem(model.topology, nonbondedMethod=app.PME, nonbondedCutoff=1.0 * u.nanometers, constraints=app.HAngles)

    integrator = mm.LangevinIntegrator(temperature, friction, timestep)

    simulation = app.Simulation(model.topology, system, integrator)
    simulation.context.setPositions(model.positions)
    print("running")
    simulation.step(1)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8
"""

core.py

Licensed under a 3-clause BSD license.

Revision history:
    12 Feb 2016 J. Sánchez-Gallego
      Initial version

"""

from __future__ import division
from __future__ import print_function
import marvin
from marvin.tools.core import MarvinUserWarning, MarvinError
from marvin.utils.db import testDbConnection
import warnings
import os

try:
    from sdss_access.path import Path
except ImportError:
    Path = None

__all__ = ['MarvinToolsClass']

mangaid_to_plateifu = {'1-209232': '8485-1901', '12-193534': '7443-3701'}

class MarvinToolsClass(object):

    def __init__(self, *args, **kwargs):
        """Marvin tools main super class.

        This super class implements the decision tree for using local files,
        database, or remote connection when initialising a Marvin tools
        object.

        """

        self.filename = kwargs.get('filename', None)
        self.mangaid = kwargs.get('mangaid', None)
        self.plateifu = kwargs.get('plateifu', None)
        self.mode = kwargs.get('mode', None)
        self.drpall = kwargs.get('drpall', None)

        if self.mode is None:
            self.mode = marvin.config.mode

        args = [self.filename, self.plateifu, self.mangaid]
        errmsg = 'Enter filename, plateifu, or mangaid!'
        assert any(args), errmsg
        assert sum([bool(arg) for arg in args]) == 1, errmsg

        if self.mangaid:
            if self.mangaid in mangaid_to_plateifu:
                self.plateifu = mangaid_to_plateifu[self.mangaid]
            else:
                raise MarvinError('mangaid={0} not found in the dictionary'
                                  .format(self.mangaid))

        if self.mode == 'local':
            self._doLocal()
        elif self.mode == 'remote':
            self._doRemote()
        elif self.mode == 'auto':
            try:
                self._doLocal()
            except:
                warnings.warn('local mode failed. Trying remote now.',
                              MarvinUserWarning)
                self._doRemote()

    def _doLocal(self):
        """Tests if it's possible to load the data locally."""

        if self.filename:

            if os.path.exists(self.filename):
                self.mode = 'local'
            else:
                raise MarvinError('input file {0} not found'
                                  .format(self.filename))

        elif self.plateifu:

            dbStatus = testDbConnection(marvin.session)
            print('db status', dbStatus)

            if dbStatus['good']:
                self.mode = 'local'
            else:
                warnings.warn(
                    'DB connection failed with error: {0}.'
                    .format(dbStatus['error']), MarvinUserWarning)

                fullpath = self._getFullPath()

                if fullpath and os.path.exists(fullpath):
                    self.mode = 'local'
                    self.filename = fullpath
                else:
                    if marvin.config.download:
                        raise NotImplementedError('sdsssync not yet implemented')
                        # When implemented, this should download the data and
                        # then kwargs['filename'] = downloaded_path and
                        # kwargs['mode'] = local
                    else:
                        raise MarvinError('this is the end of the road. Try '
                                          'using some reasonable inputs.')

    def _doRemote(self):
        """Tests if remote connection is possible."""

        if self.filename:
            raise MarvinError('filename not allowed in remote mode.')
        else:
            self.mode = 'remote'

    def _getFullPath(self, pathType, **pathParams):
        """Returns the full path of the file in the tree."""

        if not Path:
            raise MarvinError('sdss_access is not installed')
        else:
            try:
                fullpath = Path().full(pathType, **pathParams)
            except Exception as ee:
                warnings.warn(
                    'sdss_access was not able to retrieve the full path of '
                    'the file. Error message is: {0}'.format(str(ee)),
                    MarvinUserWarning)
                fullpath = None

        return fullpath
<commit_msg>replaced mangaid_to_plateifu dictionary with call to mangaid2plateifu<commit_after>#!/usr/bin/env python
# encoding: utf-8
"""

core.py

Licensed under a 3-clause BSD license.

Revision history:
    12 Feb 2016 J. Sánchez-Gallego
      Initial version

"""

from __future__ import division
from __future__ import print_function
import marvin
from marvin.tools.core import MarvinUserWarning, MarvinError
from marvin.utils.general import mangaid2plateifu
from marvin.utils.db import testDbConnection
import warnings
import os

try:
    from sdss_access.path import Path
except ImportError:
    Path = None

__all__ = ['MarvinToolsClass']

class MarvinToolsClass(object):

    def __init__(self, *args, **kwargs):
        """Marvin tools main super class.

        This super class implements the decision tree for using local files,
        database, or remote connection when initialising a Marvin tools
        object.

        """

        self.filename = kwargs.get('filename', None)
        self.mangaid = kwargs.get('mangaid', None)
        self.plateifu = kwargs.get('plateifu', None)
        self.mode = kwargs.get('mode', None)
        self.drpall = kwargs.get('drpall', None)

        if self.mode is None:
            self.mode = marvin.config.mode

        args = [self.filename, self.plateifu, self.mangaid]
        errmsg = 'Enter filename, plateifu, or mangaid!'
        assert any(args), errmsg
        assert sum([bool(arg) for arg in args]) == 1, errmsg

        if self.mangaid:
            self.plateifu = mangaid2plateifu(self.mangaid)

        if self.mode == 'local':
            self._doLocal()
        elif self.mode == 'remote':
            self._doRemote()
        elif self.mode == 'auto':
            try:
                self._doLocal()
            except:
                warnings.warn('local mode failed. Trying remote now.',
                              MarvinUserWarning)
                self._doRemote()

    def _doLocal(self):
        """Tests if it's possible to load the data locally."""

        if self.filename:

            if os.path.exists(self.filename):
                self.mode = 'local'
            else:
                raise MarvinError('input file {0} not found'
                                  .format(self.filename))

        elif self.plateifu:

            dbStatus = testDbConnection(marvin.session)
            print('db status', dbStatus)

            if dbStatus['good']:
                self.mode = 'local'
            else:
                warnings.warn(
                    'DB connection failed with error: {0}.'
                    .format(dbStatus['error']), MarvinUserWarning)

                fullpath = self._getFullPath()

                if fullpath and os.path.exists(fullpath):
                    self.mode = 'local'
                    self.filename = fullpath
                else:
                    if marvin.config.download:
                        raise NotImplementedError('sdsssync not yet implemented')
                        # When implemented, this should download the data and
                        # then kwargs['filename'] = downloaded_path and
                        # kwargs['mode'] = local
                    else:
                        raise MarvinError('this is the end of the road. Try '
                                          'using some reasonable inputs.')

    def _doRemote(self):
        """Tests if remote connection is possible."""

        if self.filename:
            raise MarvinError('filename not allowed in remote mode.')
        else:
            self.mode = 'remote'

    def _getFullPath(self, pathType, **pathParams):
        """Returns the full path of the file in the tree."""

        if not Path:
            raise MarvinError('sdss_access is not installed')
        else:
            try:
                fullpath = Path().full(pathType, **pathParams)
            except Exception as ee:
                warnings.warn(
                    'sdss_access was not able to retrieve the full path of '
                    'the file. Error message is: {0}'.format(str(ee)),
                    MarvinUserWarning)
                fullpath = None

        return fullpath
<|endoftext|>


<commit_before>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2017, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Affero General Public License for more details.

#  You should have received a copy of the GNU Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.
import collections
import logging
from h5py._hl.dataset import Dataset
from h5py._hl.group import Group
import numpy
try:
    from functools import lru_cache
except ImportError:
    from openquake.risklib.utils import memoized
else:
    memoized = lru_cache(100)
from openquake.baselib.general import DictArray
from openquake.baselib.python3compat import encode
from openquake.commonlib import calc

def extract_(dstore, dspath):
    """
    Extracts an HDF5 path object from the datastore, for instance
    extract('sitecol', dstore)
    """
    obj = dstore[dspath]
    if isinstance(obj, Dataset):
        return ArrayWrapper(obj.value, obj.attrs)
    elif isinstance(obj, Group):
        return ArrayWrapper(numpy.array(list(obj)), obj.attrs)
    else:
        return obj

class Extract(collections.OrderedDict):
    """
    A callable dictionary of functions with a single instance called
    `extract`. Then `extract(dstore, fullkey)` dispatches to the function
    determined by the first part of `fullkey` (a slash-separated
    string) by passing as argument the second part of `fullkey`.

    For instance extract(dstore, 'sitecol), extract(dstore, 'asset_values/0')
    etc.
    """
    def add(self, key, cache=False):
        def decorator(func):
            self[key] = memoized(func) if cache else func
            return func
        return decorator

    def __call__(self, dstore, key, *extra):
        try:
            k, v = key.split('/', 1)
        except ValueError:   # no slashes
            k, v = key, ''
        if k in self:
            return self[k](dstore, v, *extra)
        else:
            return extract_(dstore, key)

extract = Extract()

class ArrayWrapper(object):
    """
    A pickleable wrapper over an HDF5 dataset or group
    """
    def __init__(self, array, attrs):
        vars(self).update(attrs)
        self.array = array

    def __iter__(self):
        return iter(self.array)

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

    def __getitem__(self, idx):
        return self.array[idx]

    def __toh5__(self):
        return (self.array, {k: v for k, v in vars(self).items()
                             if k != 'array' and not k.startswith('_')})

    def __fromh5__(self, array, attrs):
        self.__init__(array, attrs)

    @property
    def dtype(self):
        return self.array.dtype

@extract.add('asset_values', cache=True)
def extract_asset_values(dstore, sid):
    """
    Extract an array of asset values for the given sid. Use it as
    /extract/asset_values/0

    :returns:
        (aid, loss_type1, ..., loss_typeN) composite array
    """
    if sid:
        return extract(dstore, 'asset_values')[int(sid)]
    asset_refs = extract(dstore, 'asset_refs')
    assetcol = extract(dstore, 'assetcol')
    assets_by_site = assetcol.assets_by_site()
    lts = assetcol.loss_types
    time_event = assetcol.time_event
    dt = numpy.dtype([('aref', asset_refs.dtype), ('aid', numpy.uint32)] +
                     [(str(lt), numpy.float32) for lt in lts])
    data = []
    for assets in assets_by_site:
        vals = numpy.zeros(len(assets), dt)
        for a, asset in enumerate(assets):
            vals[a]['aref'] = asset_refs[asset.idx]
            vals[a]['aid'] = asset.ordinal
            for lt in lts:
                vals[a][lt] = asset.value(lt, time_event)
        data.append(vals)
    return data

def convert_to_array(pmap, nsites, imtls):
    """
    Convert the probability map into a composite array with header
    of the form PGA-0.1, PGA-0.2 ...

    :param pmap: probability map
    :param nsites: total number of sites
    :param imtls: a DictArray with IMT and levels
    :returns: a composite array of lenght nsites
    """
    lst = []
    # build the export dtype, of the form PGA-0.1, PGA-0.2 ...
    for imt, imls in imtls.items():
        for iml in imls:
            lst.append(('%s-%s' % (imt, iml), numpy.float64))
    curves = numpy.zeros(nsites, numpy.dtype(lst))
    for sid, pcurve in pmap.items():
        curve = curves[sid]
        idx = 0
        for imt, imls in imtls.items():
            for iml in imls:
                curve['%s-%s' % (imt, iml)] = pcurve.array[idx]
                idx += 1
    return curves

@extract.add('hazard')
def extract_hazard(dstore, what):
    """
    Extracts hazard curves and possibly hazard maps and/or uniform hazard
    spectra. Use it as /extract/hazard/mean or /extract/hazard/rlz-0, etc
    """
    oq = dstore['oqparam']
    sitecol = dstore['sitecol']
    yield 'sitecol', sitecol
    yield 'oqparam', dstore['oqparam']
    yield 'realizations', dstore['realizations'].value
    yield 'checksum32', dstore['/'].attrs['checksum32']
    N = len(sitecol)
    if oq.poes:
        pdic = DictArray({imt: oq.poes for imt in oq.imtls})
    for kind, hcurves in calc.PmapGetter(dstore).items(what):
        logging.info('extracting hazard/%s', kind)
        yield 'hcurves-' + kind, convert_to_array(hcurves, N, oq.imtls)
        if oq.poes and oq.uniform_hazard_spectra:
            yield 'uhs-' + kind, calc.make_uhs(hcurves, oq.imtls, oq.poes, N)
        if oq.poes and oq.hazard_maps:
            hmaps = calc.make_hmap(hcurves, oq.imtls, oq.poes)
            yield 'hmaps-' + kind, convert_to_array(hmaps, N, pdic)

def _agg(losses, idxs):
    shp = losses.shape[1:]
    if not idxs:
        # no intersection, return a 0-dim matrix
        return numpy.zeros((0,) + shp, losses.dtype)
    # numpy.array wants lists, not sets, hence the sorted below
    return losses[numpy.array(sorted(idxs))].sum(axis=0)

def _filter_agg(assetcol, losses, selected):
    # losses is an array of shape (A, ..., R) with A=#assets, R=#realizations
    idxs = set(range(len(assetcol)))
    tagnames = []
    for tag in selected:
        tagname, tagvalue = tag.split('=', 1)
        if tagvalue == '*':
            tagnames.append(tagname)
        else:
            idxs &= assetcol.aids_by_tag[tag]
    if len(tagnames) > 1:
        raise ValueError('Too many * as tag values in %s' % tagnames)
    elif not tagnames:  # return an array of shape (1, ..., R)
        return ArrayWrapper(
            _agg(losses, idxs), dict(selected=encode(selected)))
    else:  # return an array of shape (T, ..., R)
        tags = [t for t in assetcol.tags() if t.startswith(tagname)]
        all_idxs = [idxs & assetcol.aids_by_tag[t] for t in tags]
        # NB: using a generator expression for all_idxs caused issues (?)
        return ArrayWrapper(
            numpy.array([_agg(losses, idxs) for idxs in all_idxs]),
            dict(selected=encode(selected), tags=encode(tags)))

@extract.add('agglosses')
def extract_agglosses(dstore, loss_type, *tags):
    """
    Aggregate losses of the given loss type and tags. Use it as
    /extract/agglosses/structural?taxonomy=RC&zipcode=20126
    /extract/agglosses/structural?taxonomy=RC&zipcode=*

    :returns:
        an array of shape (T, R) if one of the tag names has a `*` value
        an array of length 0 if there is no data for the given tags
        an array of shape (1, R), being R the number of realizations
    """
    if not loss_type:
        raise ValueError('loss_type not passed in agglosses/<loss_type>')
    l = dstore['oqparam'].lti[loss_type]
    if 'losses_by_asset' in dstore:  # scenario_risk
        losses = dstore['losses_by_asset'][:, :, l]['mean']
    elif 'avg_losses-rlzs' in dstore:  # event_based_risk
        losses = dstore['avg_losses-rlzs'][:, :, l]
    else:
        raise KeyError('No losses found in %s' % dstore)
    return _filter_agg(dstore['assetcol'], losses, tags)

@extract.add('aggdamages')
def extract_aggdamages(dstore, loss_type, *tags):
    """
    Aggregate damages of the given loss type and tags. Use it as
    /extract/aggdamages/structural?taxonomy=RC&zipcode=20126

    :returns:
        array of shape (R, D), being R the number of realizations and D
        the number of damage states or array of length 0 if there is no
        data for the given tags
    """
    if 'dmg_by_asset' in dstore:  # scenario_damage
        losses = dstore['dmg_by_asset'][loss_type]['mean']
    else:
        raise KeyError('No damages found in %s' % dstore)
    return _filter_agg(dstore['assetcol'], losses, tags)

@extract.add('aggcurves')
def extract_aggcurves(dstore, loss_type, *tags):
    """
    Aggregate loss curves of the given loss type and tags for
    event based risk calculations. Use it as
    /extract/aggcurves/structural?taxonomy=RC&zipcode=20126

    :returns:
        array of shape (S, P), being P the number of return periods
        and S the number of statistics
    """
    if 'curves-stats' in dstore:  # event_based_risk
        losses = dstore['curves-stats'][loss_type]
    else:
        raise KeyError('No curves found in %s' % dstore)
    curves = _filter_agg(dstore['assetcol'], losses, tags)
    curves.update(dstore.get_attrs('curves-stats'))
    return curves
<commit_msg>Small fix [skip hazardlib] [demos]<commit_after>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2017, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Affero General Public License for more details.

#  You should have received a copy of the GNU Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.
import collections
import logging
from h5py._hl.dataset import Dataset
from h5py._hl.group import Group
import numpy
try:
    from functools import lru_cache
except ImportError:
    from openquake.risklib.utils import memoized
else:
    memoized = lru_cache(100)
from openquake.baselib.general import DictArray
from openquake.baselib.python3compat import encode
from openquake.commonlib import calc

def extract_(dstore, dspath):
    """
    Extracts an HDF5 path object from the datastore, for instance
    extract('sitecol', dstore)
    """
    obj = dstore[dspath]
    if isinstance(obj, Dataset):
        return ArrayWrapper(obj.value, obj.attrs)
    elif isinstance(obj, Group):
        return ArrayWrapper(numpy.array(list(obj)), obj.attrs)
    else:
        return obj

class Extract(collections.OrderedDict):
    """
    A callable dictionary of functions with a single instance called
    `extract`. Then `extract(dstore, fullkey)` dispatches to the function
    determined by the first part of `fullkey` (a slash-separated
    string) by passing as argument the second part of `fullkey`.

    For instance extract(dstore, 'sitecol), extract(dstore, 'asset_values/0')
    etc.
    """
    def add(self, key, cache=False):
        def decorator(func):
            self[key] = memoized(func) if cache else func
            return func
        return decorator

    def __call__(self, dstore, key, *extra):
        try:
            k, v = key.split('/', 1)
        except ValueError:   # no slashes
            k, v = key, ''
        if k in self:
            return self[k](dstore, v, *extra)
        else:
            return extract_(dstore, key)

extract = Extract()

class ArrayWrapper(object):
    """
    A pickleable wrapper over an HDF5 dataset or group
    """
    def __init__(self, array, attrs):
        vars(self).update(attrs)
        self.array = array

    def __iter__(self):
        return iter(self.array)

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

    def __getitem__(self, idx):
        return self.array[idx]

    def __toh5__(self):
        return (self.array, {k: v for k, v in vars(self).items()
                             if k != 'array' and not k.startswith('_')})

    def __fromh5__(self, array, attrs):
        self.__init__(array, attrs)

    @property
    def dtype(self):
        return self.array.dtype

@extract.add('asset_values', cache=True)
def extract_asset_values(dstore, sid):
    """
    Extract an array of asset values for the given sid. Use it as
    /extract/asset_values/0

    :returns:
        (aid, loss_type1, ..., loss_typeN) composite array
    """
    if sid:
        return extract(dstore, 'asset_values')[int(sid)]
    asset_refs = extract(dstore, 'asset_refs')
    assetcol = extract(dstore, 'assetcol')
    assets_by_site = assetcol.assets_by_site()
    lts = assetcol.loss_types
    time_event = assetcol.time_event
    dt = numpy.dtype([('aref', asset_refs.dtype), ('aid', numpy.uint32)] +
                     [(str(lt), numpy.float32) for lt in lts])
    data = []
    for assets in assets_by_site:
        vals = numpy.zeros(len(assets), dt)
        for a, asset in enumerate(assets):
            vals[a]['aref'] = asset_refs[asset.idx]
            vals[a]['aid'] = asset.ordinal
            for lt in lts:
                vals[a][lt] = asset.value(lt, time_event)
        data.append(vals)
    return data

def convert_to_array(pmap, nsites, imtls):
    """
    Convert the probability map into a composite array with header
    of the form PGA-0.1, PGA-0.2 ...

    :param pmap: probability map
    :param nsites: total number of sites
    :param imtls: a DictArray with IMT and levels
    :returns: a composite array of lenght nsites
    """
    lst = []
    # build the export dtype, of the form PGA-0.1, PGA-0.2 ...
    for imt, imls in imtls.items():
        for iml in imls:
            lst.append(('%s-%s' % (imt, iml), numpy.float64))
    curves = numpy.zeros(nsites, numpy.dtype(lst))
    for sid, pcurve in pmap.items():
        curve = curves[sid]
        idx = 0
        for imt, imls in imtls.items():
            for iml in imls:
                curve['%s-%s' % (imt, iml)] = pcurve.array[idx]
                idx += 1
    return curves

@extract.add('hazard')
def extract_hazard(dstore, what):
    """
    Extracts hazard curves and possibly hazard maps and/or uniform hazard
    spectra. Use it as /extract/hazard/mean or /extract/hazard/rlz-0, etc
    """
    oq = dstore['oqparam']
    sitecol = dstore['sitecol']
    yield 'sitecol', sitecol
    yield 'oqparam', dstore['oqparam']
    yield 'realizations', dstore['realizations'].value
    yield 'checksum32', dstore['/'].attrs['checksum32']
    N = len(sitecol)
    if oq.poes:
        pdic = DictArray({imt: oq.poes for imt in oq.imtls})
    for kind, hcurves in calc.PmapGetter(dstore).items(what):
        logging.info('extracting hazard/%s', kind)
        yield 'hcurves-' + kind, convert_to_array(hcurves, N, oq.imtls)
        if oq.poes and oq.uniform_hazard_spectra:
            yield 'uhs-' + kind, calc.make_uhs(hcurves, oq.imtls, oq.poes, N)
        if oq.poes and oq.hazard_maps:
            hmaps = calc.make_hmap(hcurves, oq.imtls, oq.poes)
            yield 'hmaps-' + kind, convert_to_array(hmaps, N, pdic)

def _agg(losses, idxs):
    shp = losses.shape[1:]
    if not idxs:
        # no intersection, return a 0-dim matrix
        return numpy.zeros((0,) + shp, losses.dtype)
    # numpy.array wants lists, not sets, hence the sorted below
    return losses[numpy.array(sorted(idxs))].sum(axis=0)

def _filter_agg(assetcol, losses, selected):
    # losses is an array of shape (A, ..., R) with A=#assets, R=#realizations
    idxs = set(range(len(assetcol)))
    tagnames = []
    for tag in selected:
        tagname, tagvalue = tag.split('=', 1)
        if tagvalue == '*':
            tagnames.append(tagname)
        else:
            idxs &= assetcol.aids_by_tag[tag]
    if len(tagnames) > 1:
        raise ValueError('Too many * as tag values in %s' % tagnames)
    elif not tagnames:  # return an array of shape (1, ..., R)
        return ArrayWrapper(
            _agg(losses, idxs), dict(selected=encode(selected)))
    else:  # return an array of shape (T, ..., R)
        tags = [t for t in assetcol.tags() if t.startswith(tagname)]
        all_idxs = [idxs & assetcol.aids_by_tag[t] for t in tags]
        # NB: using a generator expression for all_idxs caused issues (?)
        return ArrayWrapper(
            numpy.array([_agg(losses, idxs) for idxs in all_idxs]),
            dict(selected=encode(selected), tags=encode(tags)))

@extract.add('agglosses')
def extract_agglosses(dstore, loss_type, *tags):
    """
    Aggregate losses of the given loss type and tags. Use it as
    /extract/agglosses/structural?taxonomy=RC&zipcode=20126
    /extract/agglosses/structural?taxonomy=RC&zipcode=*

    :returns:
        an array of shape (T, R) if one of the tag names has a `*` value
        an array of length 0 if there is no data for the given tags
        an array of shape (1, R), being R the number of realizations
    """
    if not loss_type:
        raise ValueError('loss_type not passed in agglosses/<loss_type>')
    l = dstore['oqparam'].lti[loss_type]
    if 'losses_by_asset' in dstore:  # scenario_risk
        losses = dstore['losses_by_asset'][:, :, l]['mean']
    elif 'avg_losses-rlzs' in dstore:  # event_based_risk
        losses = dstore['avg_losses-rlzs'][:, :, l]
    else:
        raise KeyError('No losses found in %s' % dstore)
    return _filter_agg(dstore['assetcol'], losses, tags)

@extract.add('aggdamages')
def extract_aggdamages(dstore, loss_type, *tags):
    """
    Aggregate damages of the given loss type and tags. Use it as
    /extract/aggdamages/structural?taxonomy=RC&zipcode=20126

    :returns:
        array of shape (R, D), being R the number of realizations and D
        the number of damage states or array of length 0 if there is no
        data for the given tags
    """
    if 'dmg_by_asset' in dstore:  # scenario_damage
        losses = dstore['dmg_by_asset'][loss_type]['mean']
    else:
        raise KeyError('No damages found in %s' % dstore)
    return _filter_agg(dstore['assetcol'], losses, tags)

@extract.add('aggcurves')
def extract_aggcurves(dstore, loss_type, *tags):
    """
    Aggregate loss curves of the given loss type and tags for
    event based risk calculations. Use it as
    /extract/aggcurves/structural?taxonomy=RC&zipcode=20126

    :returns:
        array of shape (S, P), being P the number of return periods
        and S the number of statistics
    """
    if 'curves-stats' in dstore:  # event_based_risk
        losses = dstore['curves-stats'][loss_type]
    else:
        raise KeyError('No curves found in %s' % dstore)
    curves = _filter_agg(dstore['assetcol'], losses, tags)
    vars(curves).update(dstore.get_attrs('curves-stats'))
    return curves
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

try:  # py3
    from shlex import quote
except ImportError:  # py2
    from pipes import quote
import logging
import os
import subprocess
import sys
import time

from threading import Thread
from getpass import getuser

from ray.autoscaler.tags import TAG_RAY_NODE_STATUS, TAG_RAY_RUNTIME_CONFIG
from ray.autoscaler.log_timer import LogTimer

logger = logging.getLogger(__name__)

# How long to wait for a node to start, in seconds
NODE_START_WAIT_S = 300
SSH_CHECK_INTERVAL = 5
CONTROL_PATH_MAX_LENGTH = 70

def get_default_ssh_options(private_key, connect_timeout, ssh_control_path):
    OPTS = [
        ("ConnectTimeout", "{}s".format(connect_timeout)),
        ("StrictHostKeyChecking", "no"),
        ("ControlMaster", "auto"),
        ("ControlPath", "{}/%C".format(ssh_control_path)),
        ("ControlPersist", "10s"),
    ]

    return ["-i", private_key] + [
        x for y in (["-o", "{}={}".format(k, v)] for k, v in OPTS) for x in y
    ]

class NodeUpdater(object):
    """A process for syncing files and running init commands on a node."""

    def __init__(self,
                 node_id,
                 provider_config,
                 provider,
                 auth_config,
                 cluster_name,
                 file_mounts,
                 initialization_commands,
                 setup_commands,
                 runtime_hash,
                 process_runner=subprocess,
                 exit_on_update_fail=False,
                 use_internal_ip=False):

        ssh_control_path = "/tmp/{}_ray_ssh_sockets/{}".format(
            getuser(), cluster_name)[:CONTROL_PATH_MAX_LENGTH]

        self.daemon = True
        self.process_runner = process_runner
        self.node_id = node_id
        self.use_internal_ip = (use_internal_ip or provider_config.get(
            "use_internal_ips", False))
        self.provider = provider
        self.ssh_private_key = auth_config["ssh_private_key"]
        self.ssh_user = auth_config["ssh_user"]
        self.ssh_control_path = ssh_control_path
        self.ssh_ip = None
        self.file_mounts = {
            remote: os.path.expanduser(local)
            for remote, local in file_mounts.items()
        }
        self.initialization_commands = initialization_commands
        self.setup_commands = setup_commands
        self.exit_on_update_fail = exit_on_update_fail
        self.runtime_hash = runtime_hash

    def get_node_ip(self):
        if self.use_internal_ip:
            return self.provider.internal_ip(self.node_id)
        else:
            return self.provider.external_ip(self.node_id)

    def wait_for_ip(self, deadline):
        while time.time() < deadline and \
                not self.provider.is_terminated(self.node_id):
            logger.info("NodeUpdater: "
                        "Waiting for IP of {}...".format(self.node_id))
            ip = self.get_node_ip()
            if ip is not None:
                return ip
            time.sleep(10)

        return None

    def set_ssh_ip_if_required(self):
        if self.ssh_ip is not None:
            return

        # We assume that this never changes.
        #   I think that's reasonable.
        deadline = time.time() + NODE_START_WAIT_S
        with LogTimer("NodeUpdater: {}: Got IP".format(self.node_id)):
            ip = self.wait_for_ip(deadline)
            assert ip is not None, "Unable to find IP of node"

        self.ssh_ip = ip

        # This should run before any SSH commands and therefore ensure that
        #   the ControlPath directory exists, allowing SSH to maintain
        #   persistent sessions later on.
        with open("/dev/null", "w") as redirect:
            try:
                self.process_runner.check_call(
                    ["mkdir", "-p", self.ssh_control_path],
                    stdout=redirect,
                    stderr=redirect)
            except subprocess.CalledProcessError as e:
                logger.warning(e)

            try:
                self.process_runner.check_call(
                    ["chmod", "0700", self.ssh_control_path],
                    stdout=redirect,
                    stderr=redirect)
            except subprocess.CalledProcessError as e:
                logger.warning(e)

    def run(self):
        logger.info("NodeUpdater: "
                    "{}: Updating to {}".format(self.node_id,
                                                self.runtime_hash))
        try:
            m = "{}: Applied config {}".format(self.node_id, self.runtime_hash)
            with LogTimer("NodeUpdater: {}".format(m)):
                self.do_update()
        except Exception as e:
            error_str = str(e)
            if hasattr(e, "cmd"):
                error_str = "(Exit Status {}) {}".format(
                    e.returncode, " ".join(e.cmd))
            logger.error("NodeUpdater: "
                         "{}: Error updating {}".format(
                             self.node_id, error_str))
            self.provider.set_node_tags(self.node_id,
                                        {TAG_RAY_NODE_STATUS: "update-failed"})
            raise e

        self.provider.set_node_tags(
            self.node_id, {
                TAG_RAY_NODE_STATUS: "up-to-date",
                TAG_RAY_RUNTIME_CONFIG: self.runtime_hash
            })

        self.exitcode = 0

    def wait_for_ssh(self, deadline):
        logger.info("NodeUpdater: "
                    "{}: Waiting for SSH...".format(self.node_id))

        while time.time() < deadline and \
                not self.provider.is_terminated(self.node_id):
            try:
                logger.debug("NodeUpdater: "
                             "{}: Waiting for SSH...".format(self.node_id))

                # Setting redirect=False allows the user to see errors like
                # unix_listener: path "/tmp/rkn_ray_ssh_sockets/..." too long
                # for Unix domain socket.
                self.ssh_cmd("uptime", connect_timeout=5, redirect=False)

                return True

            except Exception as e:
                retry_str = str(e)
                if hasattr(e, "cmd"):
                    retry_str = "(Exit Status {}): {}".format(
                        e.returncode, " ".join(e.cmd))
                logger.debug("NodeUpdater: "
                             "{}: SSH not up, retrying: {}".format(
                                 self.node_id, retry_str))
                time.sleep(SSH_CHECK_INTERVAL)

        return False

    def sync_file_mounts(self, sync_cmd):
        # Rsync file mounts
        for remote_path, local_path in self.file_mounts.items():
            assert os.path.exists(local_path), local_path
            if os.path.isdir(local_path):
                if not local_path.endswith("/"):
                    local_path += "/"
                if not remote_path.endswith("/"):
                    remote_path += "/"

            m = "{}: Synced {} to {}".format(self.node_id, local_path,
                                             remote_path)
            with LogTimer("NodeUpdater {}".format(m)):
                self.ssh_cmd(
                    "mkdir -p {}".format(os.path.dirname(remote_path)),
                    redirect=None,
                )
                sync_cmd(local_path, remote_path, redirect=None)

    def do_update(self):
        self.provider.set_node_tags(self.node_id,
                                    {TAG_RAY_NODE_STATUS: "waiting-for-ssh"})

        deadline = time.time() + NODE_START_WAIT_S
        self.set_ssh_ip_if_required()

        # Wait for SSH access
        with LogTimer("NodeUpdater: " "{}: Got SSH".format(self.node_id)):
            ssh_ok = self.wait_for_ssh(deadline)
            assert ssh_ok, "Unable to SSH to node"

        self.provider.set_node_tags(self.node_id,
                                    {TAG_RAY_NODE_STATUS: "syncing-files"})
        self.sync_file_mounts(self.rsync_up)

        # Run init commands
        self.provider.set_node_tags(self.node_id,
                                    {TAG_RAY_NODE_STATUS: "setting-up"})
        m = "{}: Initialization commands completed".format(self.node_id)
        with LogTimer("NodeUpdater: {}".format(m)):
            for cmd in self.initialization_commands:
                self.ssh_cmd(cmd, exit_on_fail=self.exit_on_update_fail)

        m = "{}: Setup commands completed".format(self.node_id)
        with LogTimer("NodeUpdater: {}".format(m)):
            for cmd in self.setup_commands:
                self.ssh_cmd(cmd, exit_on_fail=self.exit_on_update_fail)

    def rsync_up(self, source, target, redirect=None):
        logger.info("NodeUpdater: "
                    "{}: Syncing {} to {}...".format(self.node_id, source,
                                                     target))
        self.set_ssh_ip_if_required()
        self.process_runner.check_call(
            [
                "rsync", "-e", " ".join(["ssh"] + get_default_ssh_options(
                    self.ssh_private_key, 120, self.ssh_control_path)), "-avz",
                source, "{}@{}:{}".format(self.ssh_user, self.ssh_ip, target)
            ],
            stdout=redirect or sys.stdout,
            stderr=redirect or sys.stderr)

    def rsync_down(self, source, target, redirect=None):
        logger.info("NodeUpdater: "
                    "{}: Syncing {} from {}...".format(self.node_id, source,
                                                       target))
        self.set_ssh_ip_if_required()
        self.process_runner.check_call(
            [
                "rsync", "-e", " ".join(["ssh"] + get_default_ssh_options(
                    self.ssh_private_key, 120, self.ssh_control_path)), "-avz",
                "{}@{}:{}".format(self.ssh_user, self.ssh_ip, source), target
            ],
            stdout=redirect or sys.stdout,
            stderr=redirect or sys.stderr)

    def ssh_cmd(self,
                cmd,
                connect_timeout=120,
                redirect=None,
                allocate_tty=False,
                emulate_interactive=True,
                exit_on_fail=False,
                port_forward=None):

        self.set_ssh_ip_if_required()

        logger.info("NodeUpdater: Running {} on {}...".format(
            cmd, self.ssh_ip))
        ssh = ["ssh"]
        if allocate_tty:
            ssh.append("-tt")
        if emulate_interactive:
            force_interactive = (
                "true && source ~/.bashrc && "
                "export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && ")
            cmd = "bash --login -c -i {}".format(
                quote(force_interactive + cmd))

        if port_forward is None:
            ssh_opt = []
        else:
            ssh_opt = [
                "-L", "{}:localhost:{}".format(port_forward, port_forward)
            ]

        final_cmd = ssh + ssh_opt + get_default_ssh_options(
            self.ssh_private_key, connect_timeout, self.ssh_control_path) + [
                "{}@{}".format(self.ssh_user, self.ssh_ip), cmd
            ]
        try:
            self.process_runner.check_call(
                final_cmd,
                stdout=redirect or sys.stdout,
                stderr=redirect or sys.stderr)
        except subprocess.CalledProcessError:
            if exit_on_fail:
                logger.error("Command failed: \n\n  {}\n".format(
                    " ".join(final_cmd)))
                sys.exit(1)
            else:
                raise

class NodeUpdaterThread(NodeUpdater, Thread):
    def __init__(self, *args, **kwargs):
        Thread.__init__(self)
        NodeUpdater.__init__(self, *args, **kwargs)
        self.exitcode = -1
<commit_msg>[autoscaler] Fix ssh control path length issue (#5476)<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

try:  # py3
    from shlex import quote
except ImportError:  # py2
    from pipes import quote
import hashlib
import logging
import os
import subprocess
import sys
import time

from threading import Thread
from getpass import getuser

from ray.autoscaler.tags import TAG_RAY_NODE_STATUS, TAG_RAY_RUNTIME_CONFIG
from ray.autoscaler.log_timer import LogTimer

logger = logging.getLogger(__name__)

# How long to wait for a node to start, in seconds
NODE_START_WAIT_S = 300
SSH_CHECK_INTERVAL = 5
HASH_MAX_LENGTH = 10

def get_default_ssh_options(private_key, connect_timeout, ssh_control_path):
    OPTS = [
        ("ConnectTimeout", "{}s".format(connect_timeout)),
        ("StrictHostKeyChecking", "no"),
        ("ControlMaster", "auto"),
        ("ControlPath", "{}/%C".format(ssh_control_path)),
        ("ControlPersist", "10s"),
    ]

    return ["-i", private_key] + [
        x for y in (["-o", "{}={}".format(k, v)] for k, v in OPTS) for x in y
    ]

class NodeUpdater(object):
    """A process for syncing files and running init commands on a node."""

    def __init__(self,
                 node_id,
                 provider_config,
                 provider,
                 auth_config,
                 cluster_name,
                 file_mounts,
                 initialization_commands,
                 setup_commands,
                 runtime_hash,
                 process_runner=subprocess,
                 exit_on_update_fail=False,
                 use_internal_ip=False):

        ssh_control_hash = hashlib.md5(cluster_name.encode()).hexdigest()
        ssh_user_hash = hashlib.md5(getuser().encode()).hexdigest()
        ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
            ssh_user_hash[:HASH_MAX_LENGTH],
            ssh_control_hash[:HASH_MAX_LENGTH])

        self.daemon = True
        self.process_runner = process_runner
        self.node_id = node_id
        self.use_internal_ip = (use_internal_ip or provider_config.get(
            "use_internal_ips", False))
        self.provider = provider
        self.ssh_private_key = auth_config["ssh_private_key"]
        self.ssh_user = auth_config["ssh_user"]
        self.ssh_control_path = ssh_control_path
        self.ssh_ip = None
        self.file_mounts = {
            remote: os.path.expanduser(local)
            for remote, local in file_mounts.items()
        }
        self.initialization_commands = initialization_commands
        self.setup_commands = setup_commands
        self.exit_on_update_fail = exit_on_update_fail
        self.runtime_hash = runtime_hash

    def get_node_ip(self):
        if self.use_internal_ip:
            return self.provider.internal_ip(self.node_id)
        else:
            return self.provider.external_ip(self.node_id)

    def wait_for_ip(self, deadline):
        while time.time() < deadline and \
                not self.provider.is_terminated(self.node_id):
            logger.info("NodeUpdater: "
                        "Waiting for IP of {}...".format(self.node_id))
            ip = self.get_node_ip()
            if ip is not None:
                return ip
            time.sleep(10)

        return None

    def set_ssh_ip_if_required(self):
        if self.ssh_ip is not None:
            return

        # We assume that this never changes.
        #   I think that's reasonable.
        deadline = time.time() + NODE_START_WAIT_S
        with LogTimer("NodeUpdater: {}: Got IP".format(self.node_id)):
            ip = self.wait_for_ip(deadline)
            assert ip is not None, "Unable to find IP of node"

        self.ssh_ip = ip

        # This should run before any SSH commands and therefore ensure that
        #   the ControlPath directory exists, allowing SSH to maintain
        #   persistent sessions later on.
        with open("/dev/null", "w") as redirect:
            try:
                self.process_runner.check_call(
                    ["mkdir", "-p", self.ssh_control_path],
                    stdout=redirect,
                    stderr=redirect)
            except subprocess.CalledProcessError as e:
                logger.warning(e)

            try:
                self.process_runner.check_call(
                    ["chmod", "0700", self.ssh_control_path],
                    stdout=redirect,
                    stderr=redirect)
            except subprocess.CalledProcessError as e:
                logger.warning(e)

    def run(self):
        logger.info("NodeUpdater: "
                    "{}: Updating to {}".format(self.node_id,
                                                self.runtime_hash))
        try:
            m = "{}: Applied config {}".format(self.node_id, self.runtime_hash)
            with LogTimer("NodeUpdater: {}".format(m)):
                self.do_update()
        except Exception as e:
            error_str = str(e)
            if hasattr(e, "cmd"):
                error_str = "(Exit Status {}) {}".format(
                    e.returncode, " ".join(e.cmd))
            logger.error("NodeUpdater: "
                         "{}: Error updating {}".format(
                             self.node_id, error_str))
            self.provider.set_node_tags(self.node_id,
                                        {TAG_RAY_NODE_STATUS: "update-failed"})
            raise e

        self.provider.set_node_tags(
            self.node_id, {
                TAG_RAY_NODE_STATUS: "up-to-date",
                TAG_RAY_RUNTIME_CONFIG: self.runtime_hash
            })

        self.exitcode = 0

    def wait_for_ssh(self, deadline):
        logger.info("NodeUpdater: "
                    "{}: Waiting for SSH...".format(self.node_id))

        while time.time() < deadline and \
                not self.provider.is_terminated(self.node_id):
            try:
                logger.debug("NodeUpdater: "
                             "{}: Waiting for SSH...".format(self.node_id))

                # Setting redirect=False allows the user to see errors like
                # unix_listener: path "/tmp/rkn_ray_ssh_sockets/..." too long
                # for Unix domain socket.
                self.ssh_cmd("uptime", connect_timeout=5, redirect=False)

                return True

            except Exception as e:
                retry_str = str(e)
                if hasattr(e, "cmd"):
                    retry_str = "(Exit Status {}): {}".format(
                        e.returncode, " ".join(e.cmd))
                logger.debug("NodeUpdater: "
                             "{}: SSH not up, retrying: {}".format(
                                 self.node_id, retry_str))
                time.sleep(SSH_CHECK_INTERVAL)

        return False

    def sync_file_mounts(self, sync_cmd):
        # Rsync file mounts
        for remote_path, local_path in self.file_mounts.items():
            assert os.path.exists(local_path), local_path
            if os.path.isdir(local_path):
                if not local_path.endswith("/"):
                    local_path += "/"
                if not remote_path.endswith("/"):
                    remote_path += "/"

            m = "{}: Synced {} to {}".format(self.node_id, local_path,
                                             remote_path)
            with LogTimer("NodeUpdater {}".format(m)):
                self.ssh_cmd(
                    "mkdir -p {}".format(os.path.dirname(remote_path)),
                    redirect=None,
                )
                sync_cmd(local_path, remote_path, redirect=None)

    def do_update(self):
        self.provider.set_node_tags(self.node_id,
                                    {TAG_RAY_NODE_STATUS: "waiting-for-ssh"})

        deadline = time.time() + NODE_START_WAIT_S
        self.set_ssh_ip_if_required()

        # Wait for SSH access
        with LogTimer("NodeUpdater: " "{}: Got SSH".format(self.node_id)):
            ssh_ok = self.wait_for_ssh(deadline)
            assert ssh_ok, "Unable to SSH to node"

        self.provider.set_node_tags(self.node_id,
                                    {TAG_RAY_NODE_STATUS: "syncing-files"})
        self.sync_file_mounts(self.rsync_up)

        # Run init commands
        self.provider.set_node_tags(self.node_id,
                                    {TAG_RAY_NODE_STATUS: "setting-up"})
        m = "{}: Initialization commands completed".format(self.node_id)
        with LogTimer("NodeUpdater: {}".format(m)):
            for cmd in self.initialization_commands:
                self.ssh_cmd(cmd, exit_on_fail=self.exit_on_update_fail)

        m = "{}: Setup commands completed".format(self.node_id)
        with LogTimer("NodeUpdater: {}".format(m)):
            for cmd in self.setup_commands:
                self.ssh_cmd(cmd, exit_on_fail=self.exit_on_update_fail)

    def rsync_up(self, source, target, redirect=None):
        logger.info("NodeUpdater: "
                    "{}: Syncing {} to {}...".format(self.node_id, source,
                                                     target))
        self.set_ssh_ip_if_required()
        self.process_runner.check_call(
            [
                "rsync", "-e", " ".join(["ssh"] + get_default_ssh_options(
                    self.ssh_private_key, 120, self.ssh_control_path)), "-avz",
                source, "{}@{}:{}".format(self.ssh_user, self.ssh_ip, target)
            ],
            stdout=redirect or sys.stdout,
            stderr=redirect or sys.stderr)

    def rsync_down(self, source, target, redirect=None):
        logger.info("NodeUpdater: "
                    "{}: Syncing {} from {}...".format(self.node_id, source,
                                                       target))
        self.set_ssh_ip_if_required()
        self.process_runner.check_call(
            [
                "rsync", "-e", " ".join(["ssh"] + get_default_ssh_options(
                    self.ssh_private_key, 120, self.ssh_control_path)), "-avz",
                "{}@{}:{}".format(self.ssh_user, self.ssh_ip, source), target
            ],
            stdout=redirect or sys.stdout,
            stderr=redirect or sys.stderr)

    def ssh_cmd(self,
                cmd,
                connect_timeout=120,
                redirect=None,
                allocate_tty=False,
                emulate_interactive=True,
                exit_on_fail=False,
                port_forward=None):

        self.set_ssh_ip_if_required()

        logger.info("NodeUpdater: Running {} on {}...".format(
            cmd, self.ssh_ip))
        ssh = ["ssh"]
        if allocate_tty:
            ssh.append("-tt")
        if emulate_interactive:
            force_interactive = (
                "true && source ~/.bashrc && "
                "export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && ")
            cmd = "bash --login -c -i {}".format(
                quote(force_interactive + cmd))

        if port_forward is None:
            ssh_opt = []
        else:
            ssh_opt = [
                "-L", "{}:localhost:{}".format(port_forward, port_forward)
            ]

        final_cmd = ssh + ssh_opt + get_default_ssh_options(
            self.ssh_private_key, connect_timeout, self.ssh_control_path) + [
                "{}@{}".format(self.ssh_user, self.ssh_ip), cmd
            ]
        try:
            self.process_runner.check_call(
                final_cmd,
                stdout=redirect or sys.stdout,
                stderr=redirect or sys.stderr)
        except subprocess.CalledProcessError:
            if exit_on_fail:
                logger.error("Command failed: \n\n  {}\n".format(
                    " ".join(final_cmd)))
                sys.exit(1)
            else:
                raise

class NodeUpdaterThread(NodeUpdater, Thread):
    def __init__(self, *args, **kwargs):
        Thread.__init__(self)
        NodeUpdater.__init__(self, *args, **kwargs)
        self.exitcode = -1
<|endoftext|>