# 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/relax.py
# class AmberRelaxation(object):
#   def __init__(self,
#                *,
#                max_iterations: int,
#                tolerance: float,
#                stiffness: float,
#                exclude_residues: Sequence[int],
#                max_outer_iterations: int,
#                use_gpu: bool):
#   def process(self, *,
#               prot: protein.Protein) -> Tuple[str, Dict[str, Any], np.ndarray]:


import os
import numpy as np
from absl.testing import absltest
from alphafold.common import protein
from alphafold.relax import relax
#
#      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 relax."""

# Internal import (7716).


class RunAmberRelaxTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.test_dir = os.path.join(
        absltest.get_default_test_srcdir(),
        'alphafold/relax/testdata/')
    self.test_config = {
        'max_iterations': 1,
        'tolerance': 2.39,
        'stiffness': 10.0,
        'exclude_residues': [],
        'max_outer_iterations': 1,
        'use_gpu': False}

  def test_process(self):

# 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/relax.py
# class AmberRelaxation(object):
#   def __init__(self,
#                *,
#                max_iterations: int,
#                tolerance: float,
#                stiffness: float,
#                exclude_residues: Sequence[int],
#                max_outer_iterations: int,
#                use_gpu: bool):
#   def process(self, *,
#               prot: protein.Protein) -> Tuple[str, Dict[str, Any], np.ndarray]:


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

#
#      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 relax."""

# Internal import (7716).


class RunAmberRelaxTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.test_dir = os.path.join(
        absltest.get_default_test_srcdir(),
        'alphafold/relax/testdata/')
    self.test_config = {
        'max_iterations': 1,
        'tolerance': 2.39,
        'stiffness': 10.0,
        'exclude_residues': [],
        'max_outer_iterations': 1,
        'use_gpu': False}

  def test_process(self):

    amber_relax = relax.AmberRelaxation(**self.test_config)