Use numbers types and fix Parameter cache

tdgl/device/device.py

@@ -1,4 +1,5 @@
 import logging
+import numbers
 import os
 import time
 from contextlib import contextmanager, nullcontext
@@ -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

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.
@@ -155,11 +155,11 @@ def _to_tuple(items):
 
     def _evaluate(
         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]:
+    ) -> Union[Number, np.ndarray]:
         kwargs = self.kwargs.copy()
         if t is not None:
             kwargs["t"] = t
@@ -173,18 +173,21 @@ def _evaluate(
 
     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]:
+    ) -> 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:
             kwargs, kwarg_values = list(zip(*self.kwargs.items()))
@@ -279,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.
@@ -302,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, "
@@ -317,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."
             )
@@ -329,6 +331,7 @@ 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_
@@ -342,13 +345,20 @@ def __init__(
             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):
@@ -413,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:

tdgl/solution/solution.py

@@ -1,5 +1,6 @@
 import dataclasses
 import logging
+import numbers
 import operator
 import os
 import shutil
@@ -720,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()
@@ -820,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):

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:

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
@@ -216,9 +217,9 @@ def disorder_epsilon(r):
 
         # Clear the Parameter caches
         if isinstance(self.applied_vector_potential, Parameter):
-            self.applied_vector_potential._cache.clear()
+            self.applied_vector_potential._clear_cache()
         if isinstance(self.disorder_epsilon, Parameter):
-            self.disorder_epsilon._cache.clear()
+            self.disorder_epsilon._clear_cache()
 
         # Find the current terminal sites.
         self.terminal_info = device.terminal_info()
@@ -578,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,
         *,
@@ -807,9 +808,9 @@ def solve(self) -> Optional[Solution]:
 
             # Clear the Parameter caches
             if isinstance(self.applied_vector_potential, Parameter):
-                self.applied_vector_potential._cache.clear()
+                self.applied_vector_potential._clear_cache()
             if isinstance(self.disorder_epsilon, Parameter):
-                self.disorder_epsilon._cache.clear()
+                self.disorder_epsilon._clear_cache()
 
             solution = None
             if data_was_generated:

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()