Small changes to type hints

This commit is contained in:
Sebastian Kloth 2023-12-26 15:24:01 +01:00
parent eed94cfca6
commit cc185546d6
3 changed files with 8 additions and 23 deletions

View File

@ -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.

View File

@ -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]

View File

@ -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:]