Remove old score graph system and update tests (#238)

* working but flawed rework of modules tests

* move getting torsions into TorsionalEnergyNetwork

* rearrange score_support.py

* modify test_dof_space.py for removing score graph

* add ScoreSystem.intra_subscores for distinguishing poses and scoreterms

* un-comment-out test_score_graph (mistakenly commented out)

* remove some unused imports

* empty out score_module_support.py

* fix name of get_full_score_system_for parameter

* remove score imports from modules.py

* remove old graph from test_score_weights.py

* deletion of pieces of score graph and interatomic distance

* delete score_graph.py

* remove some score_graph imports from tests

* delete more remnants of score_graph

* restore delted dof minimization test

* delete some tests

* delete some score_graph imports in __init__.pys

* rename a member of cartesianenergynetwork

* move a single test from a deleted test file that we still want

* bad test_scoreterm_benchmarks.py

* first attempt at corrected score system intra total and benchmark test

* remove import

* modify test_total_gradcheck.py

* modify test_totalscore_benchmarks.py

* restore bonded path length (why was it deleted?)

* fix brackets in bonded_path_length

* fix some bonded atom imports

* remove score graph from ljlk test baseline

* update lk_ball/test_baseline.py

* update test_chemical_database.py imports

* update omega/test_baseline.py

* update rama/test_baseline.py

* fix import in test_dof_modules.py

* test_modules.py reflect intra_total namechange

* add todo notes for Tuesday

* change return value of score method intra forward

* fix ScoreSystem.intra_total

* update modules tests

* add helpers for term keywords

* modify more tests

* rename lk_ball terms

* update baseline tests

* update more tests

* fix lk_ball keys

* remove viewer

* delete intra_score_only

* delete intra_subscores

* delete old score_components.py

* delete old score_weights.py

* delete old factory_mixin.py

* delete old score_graph.py for each score method

* rename tests that used the word graph

* rename one more test parameter

* delete empty score_module_support.py

* use full word and move keyword to score_support for constraints

* update weights_keyword_to_score_method

* rename references to graph in tests

* refactor TorsionalEnergyNetwork and attempt to add mask

* rename dunbrack_one_two_three

* mask improvement and mask test

* remove unused import from test_dof_modules.py

* use specific default score weights

* Revert "remove unused import from test_dof_modules.py"

This reverts commit 8fd1562.

* masking for cartesian energy network

* change lr of LBFGS_Armijo back to 0.1 in test_modules.py

* update cartbonded/test_bench.py

* micro linting fixes

* corrections to micro linting

* remove unused variable

* add directive to ignore actually-used import

* shift noqa statement

* fix torch device in cartbonded/test_bench and ScoreSystem.intra_total

* use default torch device for minimization

* use designated device in test_scoreterm_benchmarks.py

tmol/optimization/modules.py

@@ -1,5 +1,7 @@
 import torch
-from tmol.kinematics.metadata import DOFTypes
+
+from tmol.system.kinematics import KinematicDescription
+from tmol.system.score_support import kincoords_to_coords
 
 # modules for cartesian and torsion-space optimization
 #
@@ -11,23 +13,6 @@
 #  - or we might want to keep that with dof creation
 
 
-# cartesian space minimization
-class CartesianEnergyNetwork(torch.nn.Module):
-    def __init__(self, score_graph):
-        super(CartesianEnergyNetwork, self).__init__()
-
-        # scoring graph
-        self.graph = score_graph
-
-        # parameters
-        self.dofs = torch.nn.Parameter(self.graph.coords)
-
-    def forward(self):
-        self.graph.coords = self.dofs
-        self.graph.reset_coords()
-        return self.graph.intra_score().total
-
-
 # mask out relevant dofs to the minimizer
 class DOFMaskingFunc(torch.autograd.Function):
     @staticmethod
@@ -44,25 +29,64 @@ def backward(ctx, grad_output):
         return grad, None, None
 
 
+# cartesian space minimization
+class CartesianEnergyNetwork(torch.nn.Module):
+    def __init__(self, score_system, coords, coord_mask=None):
+        super(CartesianEnergyNetwork, self).__init__()
+
+        self.score_system = score_system
+        self.coord_mask = coord_mask
+
+        self.full_coords = coords
+        if self.coord_mask is None:
+            self.masked_coords = torch.nn.Parameter(coords)
+        else:
+            self.masked_coords = torch.nn.Parameter(coords[self.coord_mask])
+
+    def forward(self):
+        self.full_coords = DOFMaskingFunc.apply(
+            self.masked_coords, self.coord_mask, self.full_coords
+        )
+        return self.score_system.intra_total(self.full_coords)
+
+
+def torsional_energy_network_from_system(score_system, residue_system, dof_mask=None):
+    # Initialize kinematic tree for the system
+    sys_kin = KinematicDescription.for_system(
+        residue_system.bonds, residue_system.torsion_metadata
+    )
+    kintree = sys_kin.kintree
+
+    # compute dofs from xyzs
+    dofs = sys_kin.extract_kincoords(residue_system.coords)
+    system_size = residue_system.system_size
+
+    return TorsionalEnergyNetwork(
+        score_system, dofs, kintree, system_size, dof_mask=dof_mask
+    )
+
+
 # torsion space minimization
 class TorsionalEnergyNetwork(torch.nn.Module):
-    def __init__(self, score_graph):
+    def __init__(self, score_system, dofs, kintree, system_size, dof_mask=None):
         super(TorsionalEnergyNetwork, self).__init__()
 
-        # scoring graph
-        self.graph = score_graph
+        self.score_system = score_system
+        self.kintree = kintree
+        self.dof_mask = dof_mask
+        self.system_size = system_size
 
-        # todo: make this a configurable parameter
-        #   (for now it defaults to torsion minimization)
-        dofmask = self.graph.dofmetadata[
-            self.graph.dofmetadata.dof_type == DOFTypes.bond_torsion
-        ]
-        self.mask = (dofmask.node_idx, dofmask.dof_idx)
+        self.full_dofs = dofs
+        if self.dof_mask is None:
+            self.masked_dofs = torch.nn.Parameter(dofs)
+        else:
+            self.masked_dofs = torch.nn.Parameter(dofs[self.dof_mask])
 
-        # parameters
-        self.dofs = torch.nn.Parameter(self.graph.dofs[self.mask])
+    def coords(self):
+        self.full_dofs = DOFMaskingFunc.apply(
+            self.masked_dofs, self.dof_mask, self.full_dofs
+        )
+        return kincoords_to_coords(self.full_dofs, self.kintree, self.system_size)
 
     def forward(self):
-        self.graph.dofs = DOFMaskingFunc.apply(self.dofs, self.mask, self.graph.dofs)
-        self.graph.reset_coords()
-        return self.graph.intra_score().total
+        return self.score_system.intra_total(self.coords())

tmol/score/__init__.py

@@ -1 +0,0 @@
-from . import viewer  # noqa: F401 import viewer to register io overloads

tmol/score/bonded_atom.py

@@ -1,6 +1,5 @@
 import attr
 
-from functools import singledispatch
 
 import torch
 import numpy
@@ -10,16 +9,9 @@
 import sparse
 import scipy.sparse.csgraph as csgraph
 
-from tmol.utility.reactive import reactive_property
 
 from tmol.types.array import NDArray
 from tmol.types.torch import Tensor
-from tmol.types.functional import validate_args
-
-from .score_graph import score_graph
-from .database import ParamDB
-from .stacked_system import StackedSystem
-from .device import TorchDevice
 
 
 @attr.s(auto_attribs=True, frozen=True, slots=True)
@@ -113,106 +105,6 @@ def to(self, device: torch.device):
         )
 
 
-@score_graph
-class BondedAtomScoreGraph(StackedSystem, ParamDB, TorchDevice):
-    """Score graph component describing a system's atom types and bonds.
-
-    Attributes:
-        atom_types: [layer, atom_index] String atom type descriptors.
-            Type descriptions defined in :py:mod:`tmol.database.chemical`.
-
-        atom_names: [layer, atom_index] String residue-specific atom name.
-
-        res_names: [layer, atom_index] String residue name descriptors.
-
-        res_indices: [layer, atom_index] Integer residue index descriptors.
-
-        bonds: [ind, (layer=0, atom_index=1, atom_index=2)] Inter-atomic bond indices.
-            Note that bonds are strictly intra-layer, and are defined by a
-            single layer index for both atoms of the bond.
-
-
-        MAX_BONDED_PATH_LENGTH: Maximum relevant inter-atomic path length.
-            Limits search depth used in ``bonded_path_length``, all longer
-            paths reported as ``inf``.
-
-    """
-
-    MAX_BONDED_PATH_LENGTH = 6
-
-    @staticmethod
-    @singledispatch
-    def factory_for(other, **_):
-        """`clone`-factory, extract atom types and bonds from other."""
-        return dict(
-            atom_types=other.atom_types,
-            atom_names=other.atom_names,
-            res_names=other.res_names,
-            res_indices=other.res_indices,
-            bonds=other.bonds,
-        )
-
-    atom_types: NDArray[object][:, :]
-    atom_names: NDArray[object][:, :]
-    res_names: NDArray[object][:, :]
-    res_indices: NDArray[int][:, :]
-    bonds: NDArray[int][:, 3]
-
-    @reactive_property
-    @validate_args
-    def real_atoms(atom_types: NDArray[object][:, :],) -> Tensor[bool][:, :]:
-        """Mask of non-null atomic indices in the system."""
-        return torch.tensor(atom_types != None)
-
-    @reactive_property
-    def indexed_bonds(bonds, system_size, device):
-        """Sorted, constant time access to bond graph."""
-        assert bonds.ndim == 2
-        assert bonds.shape[1] == 3
-
-        ## fd lkball needs this on the device
-        ibonds = IndexedBonds.from_bonds(
-            IndexedBonds.to_directed(bonds), minlength=system_size
-        ).to(device)
-
-        return ibonds
-
-    @reactive_property
-    @validate_args
-    def bonded_path_length(
-        bonds: NDArray[int][:, 3],
-        stack_depth: int,
-        system_size: int,
-        device: torch.device,
-        MAX_BONDED_PATH_LENGTH: int,
-    ) -> Tensor[float][:, :, :]:
-        """Dense inter-atomic bonded path length distance tables.
-
-        Returns:
-            [layer, from_atom, to_atom]
-            Per-layer interatomic bonded path length entries.
-        """
-
-        return torch.from_numpy(
-            bonded_path_length_stacked(
-                bonds, stack_depth, system_size, MAX_BONDED_PATH_LENGTH
-            )
-        ).to(device, dtype=torch.float)
-
-
-def bonded_path_length(
-    bonds: NDArray[int][:, 2], system_size: int, limit: int
-) -> NDArray[numpy.float32][:, :]:
-    bond_graph = sparse.COO(
-        bonds.T,
-        data=numpy.full(len(bonds), True),
-        shape=(system_size, system_size),
-        cache=True,
-    )
-
-    return csgraph.dijkstra(bond_graph, directed=False, unweighted=True, limit=limit)
-
-
 def bonded_path_length_stacked(
     bonds: NDArray[int][:, 3], stack_depth: int, system_size: int, limit: int
 ) -> NDArray[numpy.float32][:, :, :]:
@@ -230,3 +122,16 @@ def bonded_path_length_stacked(
         )
 
     return result
+
+
+def bonded_path_length(
+    bonds: NDArray[int][:, 2], system_size: int, limit: int
+) -> NDArray[numpy.float32][:, :]:
+    bond_graph = sparse.COO(
+        bonds.T,
+        data=numpy.full(len(bonds), True),
+        shape=(system_size, system_size),
+        cache=True,
+    )
+
+    return csgraph.dijkstra(bond_graph, directed=False, unweighted=True, limit=limit)

tmol/score/cartbonded/__init__.py

@@ -1 +0,0 @@
-from .score_graph import CartBondedScoreGraph  # noqa: F401