fix: simpler

Author: beckermr

jax_galsim/bounds.py

@@ -7,7 +7,6 @@
 from jax_galsim.core.utils import (
     cast_to_float,
     cast_to_int,
-    cast_to_static_numeric_scalar,
     check_is_int_then_cast,
     ensure_hashable,
     has_tracers,
@@ -509,12 +508,8 @@ def __init__(self, *args, **kwargs):
 
         self._parse_args(*args, **kwargs)
 
-        msg = (
-            "Jax-GalSim BoundsI instances must have a fixed, static width! "
-            f"Got deltax,deltay = {self.deltax!r},{self.deltay!r}."
-        )
-        self.deltax = cast_to_float(cast_to_static_numeric_scalar(self.deltax, msg=msg))
-        self.deltay = cast_to_float(cast_to_static_numeric_scalar(self.deltay, msg=msg))
+        self.deltax = cast_to_float(self.deltax)
+        self.deltay = cast_to_float(self.deltay)
         if (self.deltax != int(self.deltax)) or (self.deltay != int(self.deltay)):
             raise TypeError("BoundsI must be initialized with integer values")
         self.deltax = cast_to_int(self.deltax)

jax_galsim/core/utils.py

@@ -51,21 +51,6 @@ def has_tracers(x):
     return False
 
 
-def cast_to_static_numeric_scalar(x, msg=None):
-    if isinstance(x, (int, float, np.integer, np.floating)):
-        return x
-
-    if isinstance(x, (np.ndarray, jax.Array, jnp.ndarray)):
-        if x.ndim == 0:
-            return x.item()
-
-        if all(sv for sv in x.shape == 1):
-            return x.ravel()[0].item()
-
-    msg = msg or f"Cannot convert input {x!r} to a static, numeric scalar."
-    raise RuntimeError(msg)
-
-
 def _cast_to_type(x, typ, accept_strings=False):
     if isinstance(x, (int, float, np.integer, np.floating)) or (
         accept_strings and isinstance(x, str)

jax_galsim/photon_array.py

@@ -9,7 +9,6 @@
 from jax_galsim.core.utils import (
     cast_numpy_array_to_native_byte_order,
     cast_to_int,
-    cast_to_static_numeric_scalar,
     implements,
 )
 from jax_galsim.errors import (
@@ -92,14 +91,7 @@ def __init__(
         _nokeep=None,
     ):
         self._Ntot = _JAX_GALSIM_PHOTON_ARRAY_SIZE or N
-        self._Ntot = cast_to_int(
-            cast_to_static_numeric_scalar(
-                self._Ntot,
-                msg=(
-                    f"JAX_GalSim photon arrays must have static sizes., Got {self._Ntot!r}."
-                ),
-            )
-        )
+        self._Ntot = cast_to_int(self._Ntot)
 
         if _JAX_GALSIM_PHOTON_ARRAY_SIZE is not None and (
             self._Ntot > _JAX_GALSIM_PHOTON_ARRAY_SIZE

jax_galsim/wcs.py

@@ -1,4 +1,5 @@
 import galsim as _galsim
+import jax
 import jax.numpy as jnp
 import numpy as np
 from jax.tree_util import register_pytree_node_class
@@ -7,7 +8,6 @@
 from jax_galsim.celestial import CelestialCoord
 from jax_galsim.core.utils import (
     cast_to_float,
-    cast_to_static_numeric_scalar,
     ensure_hashable,
     implements,
 )
@@ -18,8 +18,23 @@
 from jax_galsim.transform import _Transform
 
 
+def _cast_to_static_numeric_scalar(x, msg=None):
+    if isinstance(x, (int, float, np.integer, np.floating)):
+        return x
+
+    if isinstance(x, (np.ndarray, jax.Array, jnp.ndarray)):
+        if x.ndim == 0:
+            return x.item()
+
+        if all(sv for sv in x.shape == 1):
+            return x.ravel()[0].item()
+
+    msg = msg or f"Cannot convert input {x!r} to a static, numeric scalar."
+    raise RuntimeError(msg)
+
+
 def _cast_to_python_float(x):
-    return cast_to_float(cast_to_static_numeric_scalar(x))
+    return cast_to_float(_cast_to_static_numeric_scalar(x))
 
 
 # We inherit from the reference BaseWCS and only redefine the methods that