fix mypy issues

OpenFreeEnergy · Oct 16, 2023 · b6da3f1 · b6da3f1
1 parent 91d32b9
commit b6da3f1
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/openfe/protocols/openmm_afe/equil_solvation_afe_method.py b/openfe/protocols/openmm_afe/equil_solvation_afe_method.py
@@ -176,7 +176,7 @@ def get_forward_and_reverse_analysis(self) -> dict[str, list[dict[str, Union[npt
                   fraction of data.
         """
 
-        forward_reverse = {}
+        forward_reverse: dict[str, list[dict[str, Union[npt.NDArray, unit.Quantity]]]] = {}
 
         for key in ['solvent', 'vacuum']:
             forward_reverse[key] = [
@@ -206,7 +206,7 @@ def get_overlap_matrices(self) -> dict[str, list[dict[str, npt.NDArray]]]:
               state i in state j
         """
         # Loop through and get the repeats and get the matrices
-        overlap_stats = {}
+        overlap_stats: dict[str, list[dict[str, npt.NDArray]]] = {}
 
         for key in ['solvent', 'vacuum']:
             overlap_stats[key] = [
@@ -240,7 +240,7 @@ def get_replica_transition_statistics(self) -> dict[str, list[dict[str, npt.NDAr
             * ``matrix``: The transition matrix estimate of a replica switching
               from state i to state j.
         """
-        repex_stats = {}
+        repex_stats: dict[str, list[dict[str, npt.NDArray]]] = {}
         try:
             for key in ['solvent', 'vacuum']:
                 repex_stats[key] = [
@@ -265,7 +265,7 @@ def get_replica_states(self) -> dict[str, list[npt.NDArray]]:
           the thermodynamic cycle, with lists of replica states
           timeseries for each repeat of that simulation type.
         """
-        replica_states = {}
+        replica_states: dict[str, list[npt.NDArray]] = {}
 
         for key in ['solvent', 'vacuum']:
             replicate_states[key] = [
@@ -286,7 +286,7 @@ def equilibration_iterations(self) -> dict[str, list[float]]:
           number of equilibration iterations for each repeat
           of that simulation type.
         """
-        equilibration_lengths = {}
+        equilibration_lengths: dict[str, list[float]] = {}
 
         for key in ['solvent', 'vacuum']:
             equilibration_lengths[key] = [
@@ -310,7 +310,7 @@ def production_iterations(self) -> dict[str, list[float]]:
           of production iterations for each repeat of that simulation
           type.
         """
-        production_lengths = {}
+        production_lengths: dict[str, list[float]] = {}
 
         for key in ['solvent', 'vacuum']:
             equilibration_lengths[key] = [