From cc185546d6afc93e89af460acd07e8fb1f53eafb Mon Sep 17 00:00:00 2001
From: Sebastian Kloth <sebastian.kloth@pkm.tu-darmstadt.de>
Date: Tue, 26 Dec 2023 15:24:01 +0100
Subject: [PATCH] Small changes to type hints

---
 src/mdevaluate/autosave.py     |  4 ++--
 src/mdevaluate/distribution.py |  8 ++++++--
 src/mdevaluate/utils.py        | 19 -------------------
 3 files changed, 8 insertions(+), 23 deletions(-)

diff --git a/src/mdevaluate/autosave.py b/src/mdevaluate/autosave.py
index 307c0b5..9666d5e 100644
--- a/src/mdevaluate/autosave.py
+++ b/src/mdevaluate/autosave.py
@@ -1,7 +1,7 @@
 import os
 import functools
 import inspect
-from typing import Optional, Callable
+from typing import Optional, Callable, Iterable
 
 import numpy as np
 from .checksum import checksum
@@ -142,7 +142,7 @@ def load_data(filename):
 
 
 def autosave_data(
-    nargs: int, kwargs_keys: Optional[list[str]] = None, version: Optional[int] = None
+    nargs: int, kwargs_keys: Optional[Iterable[str]] = None, version: Optional[int] = None
 ) -> Callable:
     """
     Enable autosaving of results for a function.
diff --git a/src/mdevaluate/distribution.py b/src/mdevaluate/distribution.py
index e57eb37..01d73eb 100644
--- a/src/mdevaluate/distribution.py
+++ b/src/mdevaluate/distribution.py
@@ -198,7 +198,11 @@ def distance_distribution(atoms, bins):
 def tetrahedral_order(atoms, reference_atoms=None):
     if reference_atoms is None:
         reference_atoms = atoms
-    indices = next_neighbors(reference_atoms, query_atoms=atoms, number_of_neighbors=4)
+    indices = next_neighbors(
+        reference_atoms,
+        query_atoms=atoms,
+        number_of_neighbors=4,
+    )[1]
     neighbors = reference_atoms[indices]
     neighbors_1, neighbors_2, neighbors_3, neighbors_4 = (
         neighbors[:, 0, :],
@@ -350,7 +354,7 @@ def next_neighbor_distribution(
         reference = atoms
     nn = next_neighbors(
         reference, query_atoms=atoms, number_of_neighbors=number_of_neighbors
-    )
+    )[1]
     resname_nn = reference.residue_names[nn]
     count_nn = (resname_nn == atoms.residue_names.reshape(-1, 1)).sum(axis=1)
     return np.histogram(count_nn, bins=bins, normed=normed)[0]
diff --git a/src/mdevaluate/utils.py b/src/mdevaluate/utils.py
index 4fc47fe..38334db 100644
--- a/src/mdevaluate/utils.py
+++ b/src/mdevaluate/utils.py
@@ -119,25 +119,6 @@ def filon_fourier_transformation(
     )
 
 
-def mask2indices(mask):
-    """
-    Return the selected indices of an array mask.
-    If the mask is two-dimensional, the indices will be calculated for the second axis.
-
-    Example:
-        >>> mask2indices([True, False, True, False])
-        array([0, 2])
-        >>> mask2indices([[True, True, False], [True, False, True]])
-        array([[0, 1], [0, 2]])
-    """
-    mask = np.array(mask)
-    if len(mask.shape) == 1:
-        indices = np.where(mask)
-    else:
-        indices = np.array([np.where(m) for m in mask])
-    return indices
-
-
 def superpose(x1, y1, x2, y2, N=100, damping=1.0):
     if x2[0] == 0:
         x2 = x2[1:]