diff --git a/docs/notebooks/logo.ipynb b/docs/notebooks/logo.ipynb index e0f008e..1b1c482 100644 --- a/docs/notebooks/logo.ipynb +++ b/docs/notebooks/logo.ipynb @@ -233,7 +233,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Constructing Voronoi polygons: 100%|█████| 4084/4084 [00:00<00:00, 21530.56it/s]\n" + "Constructing Voronoi polygons: 100%|███████████████████████████████████████████████████████████████████| 4084/4084 [00:00<00:00, 22641.44it/s]\n" ] }, { @@ -276,7 +276,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Simulating: 100%|██████████████████████████▉| 800/800 [01:59<00:00, 6.70tau/s ]\n" + "Simulating: 100%|████████████████████████████████████████████████████████████████████████████████████████▉| 800/800 [12:35<00:00, 1.06tau/s ]\n" ] } ], @@ -8699,7 +8699,7 @@ { "data": { "text/html": [ - "
SoftwareVersion
tdgl0.8.1; git revision 268e366 [2024-01-30]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Tue Jan 30 14:03:06 2024 PST
" + "
SoftwareVersion
tdgl0.8.1; git revision 25e009d [2024-04-10]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Wed Apr 10 14:29:46 2024 PDT
" ], "text/plain": [ "" diff --git a/docs/notebooks/mesh.ipynb b/docs/notebooks/mesh.ipynb index f7109a8..cd5b885 100644 --- a/docs/notebooks/mesh.ipynb +++ b/docs/notebooks/mesh.ipynb @@ -29,7 +29,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Constructing Voronoi polygons: 100%|███████| 525/525 [00:00<00:00, 21148.77it/s]\n" + "Constructing Voronoi polygons: 100%|█████████████████████████████████████████████████████████████████████| 525/525 [00:00<00:00, 20974.72it/s]\n" ] } ], diff --git a/docs/notebooks/polygons.ipynb b/docs/notebooks/polygons.ipynb index db26f4e..a8047b5 100644 --- a/docs/notebooks/polygons.ipynb +++ b/docs/notebooks/polygons.ipynb @@ -469,9 +469,9 @@ "name": "stderr", "output_type": "stream", "text": [ - "Constructing Voronoi polygons: 100%|███████| 459/459 [00:00<00:00, 16564.44it/s]\n", - "Constructing Voronoi polygons: 100%|███████| 459/459 [00:00<00:00, 16243.28it/s]\n", - "Constructing Voronoi polygons: 100%|███████| 459/459 [00:00<00:00, 17103.94it/s]\n" + "Constructing Voronoi polygons: 100%|█████████████████████████████████████████████████████████████████████| 459/459 [00:00<00:00, 17015.81it/s]\n", + "Constructing Voronoi polygons: 100%|█████████████████████████████████████████████████████████████████████| 459/459 [00:00<00:00, 17668.90it/s]\n", + "Constructing Voronoi polygons: 100%|█████████████████████████████████████████████████████████████████████| 459/459 [00:00<00:00, 17594.46it/s]\n" ] }, { @@ -520,7 +520,7 @@ { "data": { "text/html": [ - "
SoftwareVersion
tdgl0.8.1; git revision 268e366 [2024-01-30]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Tue Jan 30 13:57:41 2024 PST
" + "
SoftwareVersion
tdgl0.8.1; git revision 25e009d [2024-04-10]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Wed Apr 10 14:14:21 2024 PDT
" ], "text/plain": [ "" diff --git a/docs/notebooks/py-mesh.ipynb b/docs/notebooks/py-mesh.ipynb index 8f470eb..e6d67bf 100644 --- a/docs/notebooks/py-mesh.ipynb +++ b/docs/notebooks/py-mesh.ipynb @@ -184,7 +184,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Constructing Voronoi polygons: 100%|█████| 4185/4185 [00:00<00:00, 21925.67it/s]\n" + "Constructing Voronoi polygons: 100%|███████████████████████████████████████████████████████████████████| 4185/4185 [00:00<00:00, 22564.43it/s]\n" ] }, { @@ -364,7 +364,7 @@ { "data": { "text/html": [ - "
SoftwareVersion
tdgl0.8.1; git revision 268e366 [2024-01-30]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Tue Jan 30 13:57:24 2024 PST
" + "
SoftwareVersion
tdgl0.8.1; git revision 25e009d [2024-04-10]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Wed Apr 10 14:14:07 2024 PDT
" ], "text/plain": [ "" diff --git a/docs/notebooks/quickstart.ipynb b/docs/notebooks/quickstart.ipynb index 09885bb..8dcd715 100644 --- a/docs/notebooks/quickstart.ipynb +++ b/docs/notebooks/quickstart.ipynb @@ -277,7 +277,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Constructing Voronoi polygons: 100%|█████| 4679/4679 [00:00<00:00, 22554.63it/s]\n" + "Constructing Voronoi polygons: 100%|███████████████████████████████████████████████████████████████████| 4679/4679 [00:00<00:00, 21549.80it/s]\n" ] } ], @@ -368,8 +368,8 @@ "name": "stderr", "output_type": "stream", "text": [ - "Thermalizing: 100%|████████████████████████▉| 100/100 [00:10<00:00, 9.83tau/s ]\n", - "Simulating: 100%|██████████████████████████▉| 150/150 [00:25<00:00, 5.80tau/s ]\n" + "Thermalizing: 100%|██████████████████████████████████████████████████████████████████████████████████████▉| 100/100 [00:10<00:00, 9.48tau/s ]\n", + "Simulating: 100%|████████████████████████████████████████████████████████████████████████████████████████▉| 150/150 [00:27<00:00, 5.38tau/s ]\n" ] } ], @@ -7746,7 +7746,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Simulating: 100%|██████████████████████████▉| 200/200 [00:53<00:00, 3.74tau/s ]\n" + "Simulating: 100%|████████████████████████████████████████████████████████████████████████████████████████▉| 200/200 [04:49<00:00, 1.45s/tau ]\n" ] } ], @@ -12414,7 +12414,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Simulating: 100%|██████████████████████████▉| 200/200 [00:41<00:00, 4.82tau/s ]\n" + "Simulating: 100%|████████████████████████████████████████████████████████████████████████████████████████▉| 200/200 [00:42<00:00, 4.73tau/s ]\n" ] } ], @@ -27214,7 +27214,7 @@ { "data": { "text/html": [ - "
SoftwareVersion
tdgl0.8.1; git revision 268e366 [2024-01-30]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Tue Jan 30 14:29:41 2024 PST
" + "
SoftwareVersion
tdgl0.8.1; git revision 25e009d [2024-04-10]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Wed Apr 10 13:30:36 2024 PDT
" ], "text/plain": [ "" diff --git a/docs/notebooks/screening.ipynb b/docs/notebooks/screening.ipynb index 0b49e85..1aa050c 100644 --- a/docs/notebooks/screening.ipynb +++ b/docs/notebooks/screening.ipynb @@ -141,7 +141,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Constructing Voronoi polygons: 100%|█████| 2862/2862 [00:00<00:00, 22203.78it/s]\n" + "Constructing Voronoi polygons: 100%|███████████████████████████████████████████████████████████████████| 2862/2862 [00:00<00:00, 22832.46it/s]\n" ] } ], @@ -227,7 +227,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Simulating: 100%|██████████████████████████████▉| 5/5 [00:00<00:00, 10.88tau/s ]\n" + "Simulating: 100%|████████████████████████████████████████████████████████████████████████████████████████████▊| 5/5 [00:00<00:00, 10.05tau/s ]\n" ] } ], @@ -272,6 +272,13 @@ "Maximum sheet current density: 449.9 µA/µm\n" ] }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "OMP: Info #276: omp_set_nested routine deprecated, please use omp_set_max_active_levels instead.\n" + ] + }, { "name": "stdout", "output_type": "stream", @@ -334,7 +341,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "Simulating: 100%|██████████████████████████████▉| 5/5 [00:55<00:00, 11.16s/tau ]\n" + "Simulating: 100%|████████████████████████████████████████████████████████████████████████████████████████████▉| 5/5 [00:55<00:00, 11.17s/tau ]\n" ] } ], @@ -420,7 +427,7 @@ { "data": { "text/html": [ - "
SoftwareVersion
tdgl0.8.1; git revision 268e366 [2024-01-30]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Tue Jan 30 14:11:39 2024 PST
" + "
SoftwareVersion
tdgl0.8.1; git revision 25e009d [2024-04-10]
Numpy1.24.3
SciPy1.10.1
matplotlib3.7.1
cupyNone
numba0.57.1
IPython8.14.0
Python3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]
OSposix [darwin]
Number of CPUsPhysical: 10, Logical: 10
BLAS InfoOPENBLAS
Wed Apr 10 14:05:16 2024 PDT
" ], "text/plain": [ "" diff --git a/tdgl/device/device.py b/tdgl/device/device.py index b1a4bb0..adaf9a1 100644 --- a/tdgl/device/device.py +++ b/tdgl/device/device.py @@ -1,4 +1,5 @@ import logging +import numbers import os import time from contextlib import contextmanager, nullcontext @@ -88,7 +89,7 @@ def __init__( for polygon in [self.film] + self.holes: if not polygon.is_valid: - raise ValueError("Invalid Polygon: {polygon!r}.") + raise ValueError(f"Invalid Polygon: {polygon!r}.") if len(self.holes) != len(set(hole.name for hole in self.holes)): raise ValueError("All holes must have a unique name.") @@ -416,7 +417,7 @@ def scale( if not ( isinstance(origin, tuple) and len(origin) == 2 - and all(isinstance(val, (int, float)) for val in origin) + and all(isinstance(val, numbers.Real) for val in origin) ): raise TypeError("Origin must be a tuple of floats (x, y).") self._warn_if_mesh_exist("scale()") @@ -447,7 +448,7 @@ def rotate(self, degrees: float, origin: Tuple[float, float] = (0, 0)) -> "Devic if not ( isinstance(origin, tuple) and len(origin) == 2 - and all(isinstance(val, (int, float)) for val in origin) + and all(isinstance(val, numbers.Real) for val in origin) ): raise TypeError("Origin must be a tuple of floats (x, y).") self._warn_if_mesh_exist("rotate()") @@ -582,7 +583,7 @@ def _create_dimensionless_mesh( create_submesh=True, ) - def mesh_stats_dict(self) -> Dict[str, Union[int, float, str]]: + def mesh_stats_dict(self) -> Dict[str, Union[numbers.Real, str]]: """Returns a dictionary of information about the mesh.""" edge_lengths = self.edge_lengths areas = self.areas diff --git a/tdgl/parameter.py b/tdgl/parameter.py index ee2061c..5b0bb10 100644 --- a/tdgl/parameter.py +++ b/tdgl/parameter.py @@ -1,6 +1,7 @@ import hashlib import inspect import operator +from numbers import Number from typing import Callable, Optional, Union import cloudpickle @@ -67,9 +68,8 @@ class Parameter: as a function of position coordinates x, y (and optionally z and time t). Addition, subtraction, multiplication, and division - between multiple Parameters and/or real numbers (ints and floats) - is supported. The result of any of these operations is a - ``CompositeParameter`` object. + between multiple Parameters and/or numbers is supported. + The result of any of these operations is a ``CompositeParameter`` object. Args: func: A callable/function that actually calculates the parameter's value. @@ -84,9 +84,10 @@ class Parameter: kwargs: Keyword arguments for func. """ - __slots__ = ("func", "kwargs", "time_dependent", "_cache") + __slots__ = ("func", "kwargs", "time_dependent", "_cache", "_use_cache") def __init__(self, func: Callable, time_dependent: bool = False, **kwargs): + self._use_cache = kwargs.pop("use_cache", None) argspec = inspect.getfullargspec(func) args = argspec.args num_args = 2 @@ -152,26 +153,40 @@ def _to_tuple(items): + hex(hash(t)) ) + def _evaluate( + self, + x: Union[Number, np.ndarray], + y: Union[Number, np.ndarray], + z: Optional[Union[Number, np.ndarray]] = None, + t: Optional[float] = None, + ) -> Union[Number, np.ndarray]: + kwargs = self.kwargs.copy() + if t is not None: + kwargs["t"] = t + x, y = np.atleast_1d(x, y) + if z is not None: + kwargs["z"] = np.atleast_1d(z) + result = np.asarray(self.func(x, y, **kwargs)).squeeze() + if result.ndim == 0: + result = result.item() + return result + def __call__( self, - x: Union[int, float, np.ndarray], - y: Union[int, float, np.ndarray], - z: Optional[Union[int, float, np.ndarray]] = None, + x: Union[Number, np.ndarray], + y: Union[Number, np.ndarray], + z: Optional[Union[Number, np.ndarray]] = None, t: Optional[float] = None, - ) -> Union[int, float, np.ndarray]: - cache_key = self._hash_args(x, y, z, t) - if cache_key not in self._cache: - kwargs = self.kwargs.copy() - if t is not None: - kwargs["t"] = t - x, y = np.atleast_1d(x, y) - if z is not None: - kwargs["z"] = np.atleast_1d(z) - result = np.asarray(self.func(x, y, **kwargs)).squeeze() - if result.ndim == 0: - result = result.item() - self._cache[cache_key] = result - return self._cache[cache_key] + ) -> Union[Number, np.ndarray]: + if self._use_cache: + cache_key = self._hash_args(x, y, z, t) + if cache_key not in self._cache: + self._cache[cache_key] = self._evaluate(x, y, z, t) + return self._cache[cache_key] + return self._evaluate(x, y, z, t) + + def _clear_cache(self) -> None: + self._cache.clear() def _get_argspec(self) -> _FakeArgSpec: if self.kwargs: @@ -267,12 +282,11 @@ class CompositeParameter(Parameter): (i.e. it computes a scalar or vector quantity as a function of position coordinates x, y, z). A CompositeParameter object is created as a result of mathematical operations between Parameters, CompositeParameters, - and/or real numbers. + and/or numbers. Addition, subtraction, multiplication, division, and exponentiation - between Parameters, CompositeParameters and real numbers (ints and floats) - are supported. The result of any of these operations is a new - CompositeParameter object. + between ``Parameters``, ``CompositeParameters`` and numbers are supported. + The result of any of these operations is a new ``CompositeParameter`` object. Args: left: The object on the left-hand side of the operator. @@ -290,11 +304,11 @@ class CompositeParameter(Parameter): def __init__( self, - left: Union[int, float, Parameter, "CompositeParameter"], - right: Union[int, float, Parameter, "CompositeParameter"], + left: Union[Number, Parameter, "CompositeParameter"], + right: Union[Number, Parameter, "CompositeParameter"], operator_: Union[Callable, str], ): - valid_types = (int, float, complex, Parameter, CompositeParameter) + valid_types = (Number, Parameter, CompositeParameter) if not isinstance(left, valid_types): raise TypeError( f"Left must be a number, Parameter, or CompositeParameter, " @@ -305,7 +319,7 @@ def __init__( f"Right must be a number, Parameter, or CompositeParameter, " f"not {type(right)!r}." ) - if isinstance(left, (int, float)) and isinstance(right, (int, float)): + if isinstance(left, Number) and isinstance(right, Number): raise TypeError( "Either left or right must be a Parameter or CompositeParameter." ) @@ -317,22 +331,34 @@ def __init__( f"Unknown operator, {operator_!r}. " f"Valid operators are {list(self.VALID_OPERATORS)!r}." ) + self._cache = {} self.left = left self.right = right self.operator = operator_ self.time_dependent = False if isinstance(self.left, Parameter) and self.left.time_dependent: self.time_dependent = True + if self.left._use_cache is None: + self.left._use_cache = True if isinstance(self.right, Parameter) and self.right.time_dependent: self.time_dependent = True + if self.right._use_cache is None: + self.right._use_cache = True + + def _clear_cache(self) -> None: + self._cache.clear() + if isinstance(self.right._cache, Parameter): + self.right._clear_cache() + if isinstance(self.left, Parameter): + self.left._clear_cache() def __call__( self, - x: Union[int, float, np.ndarray], - y: Union[int, float, np.ndarray], - z: Optional[Union[int, float, np.ndarray]] = None, + x: Union[Number, np.ndarray], + y: Union[Number, np.ndarray], + z: Union[Number, np.ndarray, None] = None, t: Optional[float] = None, - ) -> Union[int, float, np.ndarray]: + ) -> Union[Number, np.ndarray]: kwargs = dict() if t is None else dict(t=t) values = [] for operand in (self.left, self.right): @@ -397,7 +423,7 @@ def __setstate__(self, state): class Constant(Parameter): """A Parameter whose value doesn't depend on position or time.""" - def __init__(self, value: Union[int, float, complex], dimensions: int = 2): + def __init__(self, value: Number, dimensions: int = 2): if dimensions not in (2, 3): raise ValueError(f"Dimensions must be 2 or 3, got {dimensions}.") if dimensions == 2: diff --git a/tdgl/solution/solution.py b/tdgl/solution/solution.py index 1c62aaf..dc11024 100644 --- a/tdgl/solution/solution.py +++ b/tdgl/solution/solution.py @@ -1,8 +1,8 @@ import dataclasses import logging +import numbers import operator import os -import pickle import shutil from contextlib import nullcontext from datetime import datetime @@ -721,7 +721,7 @@ def field_at_position( ) zs = positions[:, 2] positions = positions[:, :2] - elif isinstance(zs, (int, float, np.generic)): + elif isinstance(zs, numbers.Real): # constant zs zs = zs * np.ones(len(positions)) zs = zs.squeeze() @@ -821,7 +821,7 @@ def vector_potential_at_position( ) zs = positions[:, 2] positions = positions[:, :2] - elif isinstance(zs, (int, float, np.generic)): + elif isinstance(zs, numbers.Real): # constant zs zs = zs * np.ones(len(positions)) if not isinstance(zs, np.ndarray): @@ -970,7 +970,7 @@ def deserialize_func(name, h5group): if name in h5group.attrs: return h5group.attrs[name] if f"{name}.pickle" in h5group: - return pickle.loads(np.void(grp[f"{name}.pickle"]).tobytes()) + return cloudpickle.loads(np.void(grp[f"{name}.pickle"]).tobytes()) raise IOError(f"Unable to load {name}.") with h5py.File(path, "r") as f: diff --git a/tdgl/solver/runner.py b/tdgl/solver/runner.py index 08f6a8a..6515e0e 100644 --- a/tdgl/solver/runner.py +++ b/tdgl/solver/runner.py @@ -1,5 +1,6 @@ import itertools import logging +import numbers import os import subprocess import sys @@ -153,7 +154,7 @@ def save_fixed_values(self, fixed_data: Dict[str, np.ndarray]) -> None: def save_time_step( self, - state: Dict[str, Union[int, float]], + state: Dict[str, numbers.Real], data: Dict[str, np.ndarray], running_state: Union[Dict[str, np.ndarray], None], ) -> None: diff --git a/tdgl/solver/solver.py b/tdgl/solver/solver.py index 022d8c7..d855e44 100644 --- a/tdgl/solver/solver.py +++ b/tdgl/solver/solver.py @@ -2,6 +2,7 @@ import itertools import logging import math +import numbers import os from datetime import datetime from typing import Callable, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union @@ -214,6 +215,12 @@ def disorder_epsilon(r): if np.any(epsilon > 1): raise ValueError("The disorder parameter epsilon must be <= 1") + # Clear the Parameter caches + if isinstance(self.applied_vector_potential, Parameter): + self.applied_vector_potential._clear_cache() + if isinstance(self.disorder_epsilon, Parameter): + self.disorder_epsilon._clear_cache() + # Find the current terminal sites. self.terminal_info = device.terminal_info() self.terminal_names = [term.name for term in self.terminal_info] @@ -572,7 +579,7 @@ def get_induced_vector_potential( def update( self, - state: Dict[str, Union[int, float]], + state: Dict[str, numbers.Real], running_state: RunningState, dt: float, *, @@ -799,6 +806,12 @@ def solve(self) -> Optional[Solution]: logger.info(f"Simulation ended at {end_time}") logger.info(f"Simulation took {end_time - start_time}") + # Clear the Parameter caches + if isinstance(self.applied_vector_potential, Parameter): + self.applied_vector_potential._clear_cache() + if isinstance(self.disorder_epsilon, Parameter): + self.disorder_epsilon._clear_cache() + solution = None if data_was_generated: solution = Solution( diff --git a/tdgl/version.py b/tdgl/version.py index 25223ef..f5f4a9b 100644 --- a/tdgl/version.py +++ b/tdgl/version.py @@ -1,4 +1,4 @@ -__version_info__ = (0, 8, 2) +__version_info__ = (0, 8, 1) __version__ = ".".join(map(str, __version_info__)) diff --git a/tdgl/visualization/snapshot.py b/tdgl/visualization/snapshot.py index e37f7e3..1f875fa 100644 --- a/tdgl/visualization/snapshot.py +++ b/tdgl/visualization/snapshot.py @@ -1,3 +1,4 @@ +import numbers from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Union import h5py @@ -48,7 +49,7 @@ def generate_snapshots( Returns: The matplotlib figure and axes for each time in ``times`` """ - if isinstance(times, (int, float)): + if isinstance(times, numbers.Real): times = [times] if quantities is None: quantities = Quantity.get_keys()