# Path: alphafold/common/protein.py
# PDB_CHAIN_IDS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'
# PDB_MAX_CHAINS = len(PDB_CHAIN_IDS)  # := 62.
# class Protein:
#   def __post_init__(self):
# def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein:
# def _chain_end(atom_index, end_resname, chain_name, residue_index) -> str:
# def to_pdb(prot: Protein) -> str:
# def ideal_atom_mask(prot: Protein) -> np.ndarray:
# def from_prediction(
#     features: FeatureDict,
#     result: ModelOutput,
#     b_factors: Optional[np.ndarray] = None,
#     remove_leading_feature_dimension: bool = True) -> Protein:
#   def _maybe_remove_leading_dim(arr: np.ndarray) -> np.ndarray:
# 
# Path: alphafold/relax/amber_minimize.py
# ENERGY = unit.kilocalories_per_mole
# LENGTH = unit.angstroms
# def will_restrain(atom: openmm_app.Atom, rset: str) -> bool:
# def _add_restraints(
#     system: openmm.System,
#     reference_pdb: openmm_app.PDBFile,
#     stiffness: unit.Unit,
#     rset: str,
#     exclude_residues: Sequence[int]):
# def _openmm_minimize(
#     pdb_str: str,
#     max_iterations: int,
#     tolerance: unit.Unit,
#     stiffness: unit.Unit,
#     restraint_set: str,
#     exclude_residues: Sequence[int],
#     use_gpu: bool):
# def _get_pdb_string(topology: openmm_app.Topology, positions: unit.Quantity):
# def _check_cleaned_atoms(pdb_cleaned_string: str, pdb_ref_string: str):
# def _check_residues_are_well_defined(prot: protein.Protein):
# def _check_atom_mask_is_ideal(prot):
# def clean_protein(
#     prot: protein.Protein,
#     checks: bool = True):
# def make_atom14_positions(prot):
# def find_violations(prot_np: protein.Protein):
# def get_violation_metrics(prot: protein.Protein):
# def _run_one_iteration(
#     *,
#     pdb_string: str,
#     max_iterations: int,
#     tolerance: float,
#     stiffness: float,
#     restraint_set: str,
#     max_attempts: int,
#     use_gpu: bool,
#     exclude_residues: Optional[Collection[int]] = None):
# def run_pipeline(
#     prot: protein.Protein,
#     stiffness: float,
#     use_gpu: bool,
#     max_outer_iterations: int = 1,
#     place_hydrogens_every_iteration: bool = True,
#     max_iterations: int = 0,
#     tolerance: float = 2.39,
#     restraint_set: str = "non_hydrogen",
#     max_attempts: int = 100,
#     checks: bool = True,
#     exclude_residues: Optional[Sequence[int]] = None):
# def get_initial_energies(pdb_strs: Sequence[str],
#                          stiffness: float = 0.0,
#                          restraint_set: str = "non_hydrogen",
#                          exclude_residues: Optional[Sequence[int]] = None):


import os
import numpy as np
from absl.testing import absltest
from alphafold.common import protein
from alphafold.relax import amber_minimize
# Copyright 2021 DeepMind Technologies Limited
#
# 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.

"""Tests for amber_minimize."""

# Internal import (7716).

_USE_GPU = False


def _load_test_protein(data_path):
  pdb_path = os.path.join(absltest.get_default_test_srcdir(), data_path)
  with open(pdb_path, 'r') as f:

# Path: alphafold/common/protein.py
# PDB_CHAIN_IDS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'
# PDB_MAX_CHAINS = len(PDB_CHAIN_IDS)  # := 62.
# class Protein:
#   def __post_init__(self):
# def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein:
# def _chain_end(atom_index, end_resname, chain_name, residue_index) -> str:
# def to_pdb(prot: Protein) -> str:
# def ideal_atom_mask(prot: Protein) -> np.ndarray:
# def from_prediction(
#     features: FeatureDict,
#     result: ModelOutput,
#     b_factors: Optional[np.ndarray] = None,
#     remove_leading_feature_dimension: bool = True) -> Protein:
#   def _maybe_remove_leading_dim(arr: np.ndarray) -> np.ndarray:
# 
# Path: alphafold/relax/amber_minimize.py
# ENERGY = unit.kilocalories_per_mole
# LENGTH = unit.angstroms
# def will_restrain(atom: openmm_app.Atom, rset: str) -> bool:
# def _add_restraints(
#     system: openmm.System,
#     reference_pdb: openmm_app.PDBFile,
#     stiffness: unit.Unit,
#     rset: str,
#     exclude_residues: Sequence[int]):
# def _openmm_minimize(
#     pdb_str: str,
#     max_iterations: int,
#     tolerance: unit.Unit,
#     stiffness: unit.Unit,
#     restraint_set: str,
#     exclude_residues: Sequence[int],
#     use_gpu: bool):
# def _get_pdb_string(topology: openmm_app.Topology, positions: unit.Quantity):
# def _check_cleaned_atoms(pdb_cleaned_string: str, pdb_ref_string: str):
# def _check_residues_are_well_defined(prot: protein.Protein):
# def _check_atom_mask_is_ideal(prot):
# def clean_protein(
#     prot: protein.Protein,
#     checks: bool = True):
# def make_atom14_positions(prot):
# def find_violations(prot_np: protein.Protein):
# def get_violation_metrics(prot: protein.Protein):
# def _run_one_iteration(
#     *,
#     pdb_string: str,
#     max_iterations: int,
#     tolerance: float,
#     stiffness: float,
#     restraint_set: str,
#     max_attempts: int,
#     use_gpu: bool,
#     exclude_residues: Optional[Collection[int]] = None):
# def run_pipeline(
#     prot: protein.Protein,
#     stiffness: float,
#     use_gpu: bool,
#     max_outer_iterations: int = 1,
#     place_hydrogens_every_iteration: bool = True,
#     max_iterations: int = 0,
#     tolerance: float = 2.39,
#     restraint_set: str = "non_hydrogen",
#     max_attempts: int = 100,
#     checks: bool = True,
#     exclude_residues: Optional[Sequence[int]] = None):
# def get_initial_energies(pdb_strs: Sequence[str],
#                          stiffness: float = 0.0,
#                          restraint_set: str = "non_hydrogen",
#                          exclude_residues: Optional[Sequence[int]] = None):


# Path: alphafold/relax/amber_minimize_test.py
import os
import numpy as np
from absl.testing import absltest
from alphafold.common import protein
from alphafold.relax import amber_minimize

# Copyright 2021 DeepMind Technologies Limited
#
# 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.

"""Tests for amber_minimize."""

# Internal import (7716).

_USE_GPU = False


def _load_test_protein(data_path):
  pdb_path = os.path.join(absltest.get_default_test_srcdir(), data_path)
  with open(pdb_path, 'r') as f:

    return protein.from_pdb_string(f.read())