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 os
import functools import functools
import inspect import inspect
from typing import Optional, Callable from typing import Optional, Callable, Iterable
import numpy as np import numpy as np
from .checksum import checksum from .checksum import checksum
@ -142,7 +142,7 @@ def load_data(filename):
def autosave_data( 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: ) -> Callable:
""" """
Enable autosaving of results for a function. 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): def tetrahedral_order(atoms, reference_atoms=None):
if reference_atoms is None: if reference_atoms is None:
reference_atoms = atoms 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 = reference_atoms[indices]
neighbors_1, neighbors_2, neighbors_3, neighbors_4 = ( neighbors_1, neighbors_2, neighbors_3, neighbors_4 = (
neighbors[:, 0, :], neighbors[:, 0, :],
@ -350,7 +354,7 @@ def next_neighbor_distribution(
reference = atoms reference = atoms
nn = next_neighbors( nn = next_neighbors(
reference, query_atoms=atoms, number_of_neighbors=number_of_neighbors reference, query_atoms=atoms, number_of_neighbors=number_of_neighbors
) )[1]
resname_nn = reference.residue_names[nn] resname_nn = reference.residue_names[nn]
count_nn = (resname_nn == atoms.residue_names.reshape(-1, 1)).sum(axis=1) count_nn = (resname_nn == atoms.residue_names.reshape(-1, 1)).sum(axis=1)
return np.histogram(count_nn, bins=bins, normed=normed)[0] 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): def superpose(x1, y1, x2, y2, N=100, damping=1.0):
if x2[0] == 0: if x2[0] == 0:
x2 = x2[1:] x2 = x2[1:]