miepom 0.2.0__py3-none-any.whl

miepom/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

miepom/backend.py

@@ -0,0 +1,33 @@
+"""Backend switching for the MIEPOM library.
+
+Backend switching allows the user to change the underlying computational library (e.g., NumPy, CuPy) used by the MIEPOM library.
+Use the `set_backend(xp)` function to switch the backend, where `xp` is the desired computational library (e.g., `numpy`, `cupy`).
+"""
+
+from .sim import set_backend as _set_simulator_backend
+from .linalg import set_backend as _set_linalg_backend
+from .sensor import set_backend as _set_sensor_backend
+from .pom import set_backend as _set_microscope_backend
+from .config import set_backend as _set_config_backend
+
+import numpy as np
+
+globals()["xp"] = np
+
+def set_backend(xp):
+    """Set the linear algebra backend for the MIEPOM library."""
+    _set_linalg_backend(xp)
+    _set_sensor_backend(xp)
+    _set_microscope_backend(xp)
+    _set_simulator_backend(xp)
+    _set_config_backend(xp)
+    globals()["xp"] = xp
+    print(f"MIEPOM backend set to {xp.__name__}")
+
+def get_backend():
+    """Get the name of the current linear algebra backend for the MIEPOM library."""
+    return globals()["xp"]
+
+def get_backend_name():
+    """Get the name of the current linear algebra backend for the MIEPOM library."""
+    return globals()["xp"].__name__

miepom/config.py

@@ -0,0 +1,216 @@
+"""Configuration settings for the MIEPOM library."""
+
+import numpy as np
+import os, time
+from .utils import print_line
+
+globals()["xp"] = np
+
+def set_backend(xp):
+    globals()["xp"] = xp
+try:
+    import scipy as sp
+    USESCIPY = True
+except ImportError:
+    USESCIPY = False
+try:
+    import numba as nb
+    USENUMBA = True
+except ImportError:
+    USENUMBA = False
+
+try:
+    import mlx.core as mx
+    USEMLX = True
+except ImportError:
+    USEMLX = False
+
+try:
+    import cupy as cp
+    USECUPY = True
+except ImportError:
+    USECUPY = False
+
+VERBOSE = False if os.environ.get("MIEPOM_VERBOSE", "0") == "0" else True
+
+# always use numpy as the default backend for the core object, but allow switching to other backends if available
+DEFAULT_BACKEND = "numpy"
+
+#: Computation precision. Note that if MLX (or CuPy) is used, it is always single precision
+PRECISION = "single" if os.environ.get("MIEPOM_USE_DOUBLE", "0") == "0" else "double" # 'single' or 'double'
+
+#: whether to use fast math (numba)
+NBFASTMATH = False if os.environ.get("MIEPOM_FASTMATH", "0") == "0" else True
+
+#: numba's compile target argument for vectorize and guvectorize functions.
+NBTARGET = os.environ.get("MIEPOM_TARGET", "parallel")
+if NBTARGET not in ["cpu", "parallel", "cuda"]:
+    raise ValueError(f"Invalid MIEPOM_TARGET value: {NBTARGET}. Must be one of 'cpu', 'parallel', or 'cuda'.")
+
+#: if set to True, it assumes a vectorized version of the miepython library
+VECTORIZED_MIE = False if os.environ.get("MIEPOM_VECTORIZED_MIE", "0") == "0" else True
+
+# Computed constants below. Do not modify!
+#-----------------------------------------
+
+if USENUMBA:
+    import numba as nb
+
+class nbtype():
+    def __getitem__(self, key):
+        pass
+
+try:
+    import mlx.core as mx
+    mx_compile = mx.compile
+    mx_eval = mx.eval
+except ImportError:
+    mx_compile = lambda f: f
+    mx_eval = lambda *args, **kwargs: None
+
+def mx_compile_out(func):
+    """A decorator to compile a function with mlx.core.compile. This is used to compile the sensor_image function with mlx.
+    It also handles the 'out' keyword argument for in-place computation.
+    """
+    func = mx_compile(func)
+    def wrapper(*args, **kwargs):
+        if 'out' in kwargs:
+            out = kwargs.pop('out')
+        else:
+            out = None
+        if out is not None:
+            out[:] = func(*args, **kwargs)[:]
+        else:
+            out = func(*args, **kwargs)
+        return out
+    return wrapper
+
+def mx_out(func):
+    def wrapper(*args, **kwargs):
+        if 'out' in kwargs:
+            out = kwargs.pop('out')
+        else:
+            out = None
+        if out is not None:
+            out[:] = func(*args, **kwargs)[:]
+        else:
+            out = func(*args, **kwargs)
+        return out
+    return wrapper
+
+if USENUMBA:
+    vectorize = nb.vectorize
+    guvectorize = nb.guvectorize
+    jit = nb.jit
+    njit = nb.njit
+else:
+    vectorize = lambda f,*args,**kwargs: f
+    guvectorize = lambda f,*args,**kwargs: f
+    jit = lambda f,*args,**kwargs: f
+    njit = lambda f,*args,**kwargs: f
+    NBFLOAT = nbtype()
+    NBCOMPLEX = nbtype()
+    NBUINT = nbtype()
+    NBBOOL = nbtype()
+
+    nbfloat = nbtype()
+    nbcomplex = nbtype()
+    nbuint = nbtype()
+    nbbool = nbtype()
+
+os.environ["MIEPYTHON_CACHE"] = "0"
+os.environ["MIEPYTHON_FASTMATH"] = "1"
+os.environ["MIEPYTHON_TARGET"] = NBTARGET
+
+if PRECISION == "single":
+    os.environ["MIEPYTHON_USE_DOUBLE"] = "0"
+    if USENUMBA:
+        NBFLOAT = nb.float32
+        NBCOMPLEX = nb.complex64
+        NBUINT = nb.uint16
+        NBBOOL = nb.boolean
+
+        nbfloat = nb.float32
+        nbcomplex = nb.complex64
+        nbuint = nb.uint16
+        nbbool = nb.boolean
+
+    CDTYPE = np.complex64
+    FDTYPE = np.float32
+    UDTYPE = "uint16"
+
+    cdtype = lambda: xp.complex64
+    fdtype = lambda: xp.float32
+    udtype = lambda: xp.uint16
+
+elif PRECISION == "double":
+    os.environ["MIEPYTHON_USE_DOUBLE"] = "1"
+    if USENUMBA:
+        NBFLOAT = nb.float64
+        NBCOMPLEX = nb.complex128
+        NBUINT = nb.uint16
+        NBBOOL = nb.boolean
+
+        nbfloat = nb.float64
+        nbcomplex = nb.complex128
+        nbuint = nb.uint16
+        nbbool = nb.boolean
+
+    CDTYPE = np.complex128
+    FDTYPE = np.float64
+    UDTYPE = "uint16"
+
+    cdtype = lambda: xp.complex128
+    fdtype = lambda: xp.float64
+    udtype = lambda: xp.uint16
+else:
+    raise ValueError("Invalid precision name")
+
+import miepython as mie
+
+def print_settings():
+    print("Library settings")
+    print_line("-")
+    
+    print(f"Using Numba:         {USENUMBA}")
+    if USENUMBA:
+        print(f"Using fastmath:      {NBFASTMATH}")
+        print(f"Numba target:        '{NBTARGET}'")
+        print(f"Num threads:         {nb.get_num_threads()}")
+    print(f"Precission:          '{PRECISION}'")
+    print(f"Using mlx:            {USEMLX}")
+    print(f"Using cupy:           {USECUPY}")
+    print_line("=")
+
+if VERBOSE:
+    print_settings()
+
+    print("numpy", np.__version__)
+    print("miepython", mie.__version__)
+    if USESCIPY:
+        print("scipy", sp.__version__)
+    if USECUPY:
+        print("cupy", cp.__version__)
+    if USEMLX:
+        print("mlx", mx.__version__)   
+    if USENUMBA:
+        print("numba", nb.__version__)
+
+    print("-----------------------------------------------------------------------")
+
+def get_array(x):
+    """Convert array to numpy array if it is a cupy or mlx array, otherwise return the input as is. 
+    This is useful for converting arrays to numpy for further processing or visualization."""
+    if xp.__name__ == 'mlx.core':
+        # Evaluate the eoutput using mlx.eval to ensure it is a computed array in the mlx backend. 
+        # This is necessary because mlx uses lazy evaluation, and we want to ensure that the result is a concrete array.
+        mx_eval(x)
+        return np.asarray(x)
+    elif xp.__name__ == 'cupy':
+        return x.get()  # Convert CuPy array to NumPy array
+    else:
+        #try conversion to numpy array, if it fails, return the input as is
+        try:
+            return np.asarray(x)
+        except Exception:
+            return x

miepom/core.py

@@ -0,0 +1,12 @@
+"""Core module for the MIEPOM library.
+
+This module provides convenient access to the main components of the MIEPOM library, including particles, simulators, microscopes, sensors, and backend management.
+"""
+
+from .particles import *
+from .sim import *
+from .pom import *
+from .sensor import  *
+from .backend import *
+
+from .config import get_array

miepom/linalg.py

@@ -0,0 +1,195 @@
+"""Linear algebra utilities for the PPOM library."""
+
+# Compiled numba code
+from .config import USENUMBA, nbfloat,nbcomplex, nbuint, nbbool, NBFASTMATH, NBTARGET, vectorize, guvectorize, jit, njit, mx_compile, mx_compile_out, mx_compile
+import numpy as np
+
+# Set the default linear algebra backend to numpy so that we can compile and run numba functions correctly.
+globals()["xp"] = np
+
+def backend_name():
+    """Return the name of the current linear algebra backend."""
+    return globals()["xp"].__name__
+
+def set_backend(xp):
+    globals()["xp"] = xp
+
+# data packing and unpacking functions for boolean masks
+
+# reference implementation of the pack function using numpy advanced indexing. 
+def _pack_np(data,mask):
+    return data[...,mask]
+
+# numba implementation of the pack function, requires a dummy argument so that we can determine the output shape at compile time. This is a hack, but it works.
+@guvectorize([(nbcomplex[:,:],nbbool[:,:],nbcomplex[:],nbcomplex[:])],
+                 "(m,n),(m,n),(l)->(l)",boundscheck = False, target = NBTARGET, fastmath = NBFASTMATH)
+def __pack_nb(a, mask, _dummy, out):
+    n = 0
+    for i in range(mask.shape[0]):
+        for j in range(mask.shape[1]):
+            if mask[i,j] == True:
+                if n < len(out):
+                    out[n] = a[i,j]
+                    n+=1
+
+# numba wrapper to remove the dummy argument
+def _pack_nb(a,mask, out = None):
+    if out is None:
+        out = np.empty(a.shape[:-2] + (np.sum(mask),), dtype = a.dtype)
+    return __pack_nb(a,mask,out,out = out) 
+
+# generic xp implementation that calls the numpy version and adds support for output array
+def _pack_xp(data,mask, out = None):
+    if out is None:
+        return _pack_np(data,mask)
+    else:   
+        out[:] = _pack_np(data,mask)
+        return out
+
+# boolean indices are not supported in mlx, so we need to implement a custom pack function for mlx
+def _pack_mx(data,mask, out = None):
+    if out is None:
+        #we convert to numpy array first, because mlx does not support boolean indexing
+        return xp.array(np.asarray(data)[...,mask])
+    else:
+        out[:] = np.asarray(data)[...,mask]
+        return out
+
+def _unpack_np(mask,data, out = None):
+    if out is None:
+        out = xp.zeros(data.shape[:-1] + mask.shape, dtype = data.dtype)
+    out[...] = 0.
+    out[...,mask] = data
+    return out
+
+_unpack_xp = _unpack_np
+
+# the mx version cannot use vectorize version because it does not support boolean indexing, so we need to implement a custom unpack function for mlx
+def __unpack_mx(mask,data, out = None):
+    if out is None:
+        out = xp.zeros(mask.shape, dtype = data.dtype)
+    out[:] = 0.
+    out[mask] = data
+    return out
+
+# the vectorized version of the mlx unpack function is used when the data is 2D, so that we can use vmap to speed up the computation.
+def __unpack_mx_vec(mask,data, out = None):
+    if out is None:
+        return xp.vmap(__unpack_mx, in_axes=(None, 0), out_axes=0)(mask,data)
+    else:
+        out[:] = 0.
+        tmp = xp.vmap(__unpack_mx, in_axes=(None, 0), out_axes=0)(mask,data)
+        out[:] = tmp
+        return out
+
+def _unpack_mx(mask,data, out = None):
+    if data.ndim == 2:
+        return __unpack_mx_vec(mask,data, out = out)
+    else:   
+        return __unpack_mx(mask,data, out = out)
+
+@jit([(nbbool[:,:],nbcomplex[:],nbcomplex[:,:])])
+def __unpack_nb(mask,data,out):
+    k = 0
+    for i in range(out.shape[0]):
+        for j in range(out.shape[1]):
+            if mask[i,j] == True:
+                out[i,j] = data[k]
+                k+=1
+            else:
+                out[i,j] = 0.
+
+@guvectorize([(nbbool[:,:],nbcomplex[:],nbcomplex[:,:])], "(n,m),(k)->(n,m)", target = NBTARGET, fastmath = NBFASTMATH)
+def _unpack_nb(mask,data,out):
+    __unpack_nb(mask,data,out)
+
+def unpack(mask, data, out = None):
+    """Unpack a 1D array into a 2D array at the locations specified by a boolean mask."""
+    if backend_name() == "mlx.core":   
+        return _unpack_mx(mask,data, out = out)
+    elif backend_name() == "numba":
+        return _unpack_nb(mask,data, out = out)
+    else:
+        return _unpack_xp(mask,data, out)
+
+def pack(data, mask, out = None):
+    """Pack a 2D array into a 1D array at the locations specified by a boolean mask."""
+    if backend_name() == "mlx.core":
+        return _pack_mx(data, mask, out)
+    elif backend_name() == "numba":
+        return _pack_nb(data, mask, out)
+    else:
+        return _pack_xp(data, mask, out)
+
+def _abs2_np(a):
+    return a.real**2 + a.imag**2
+
+_abs2_nb = vectorize([nbfloat(nbcomplex)],target = NBTARGET,fastmath = NBFASTMATH)(_abs2_np)
+
+def _abs2_mx(a, out = None):
+    if out is None:
+        return mx_compile(_abs2_np)(a)
+    else:
+        out[:] = mx_compile(_abs2_np)(a)
+        return out
+
+def _abs2_xp(a, out = None):
+    x2 = np.real(a)**2
+    y = np.imag(a)
+    y2 = np.multiply(y, y, out = out)
+    return np.add(x2, y2, out = y2)
+
+def abs2(a, out = None):
+    if backend_name() == "mlx.core":
+        return _abs2_mx(a, out)
+    elif backend_name() == "numba":
+        return _abs2_nb(a, out)
+    else:
+        return _abs2_xp(a, out)
+
+_vecdot_xp = np.vecdot
+
+def _vecdot_mx(a,b, out = None):
+    if out is None:
+        return mx_compile(_vecdot_xp)(a,b)
+    else:
+        out[:] = mx_compile(_vecdot_xp)(a,b)
+        return out  
+
+@guvectorize([(nbfloat[:],nbfloat[:],nbfloat[:])], "(n),(n)->()", 
+                fastmath = NBFASTMATH, cache = True, target = NBTARGET)
+def _vecdot_nb(a,b, out):
+    tmp = 0
+    for i in range(3):
+        tmp += a[i] * b[i]
+    out[0] = tmp
+
+def vecdot(a,b, out = None):
+    if backend_name() == "mlx.core":
+        return _vecdot_mx(a,b, out)
+    elif backend_name() == "numba":
+        return _vecdot_nb(a,b, out)
+    else:
+        return _vecdot_xp(a,b, out)
+
+def _exp_xp(x):
+    return xp.exp(x)
+
+def _exp_mx(x, out = None):
+    if out is None:
+        return mx_compile(_exp_xp)(x)
+    else:
+        out[:] = mx_compile(_exp_xp)(x)
+        return out
+
+_exp_nb = vectorize([nbcomplex(nbcomplex)], target = NBTARGET, fastmath = NBFASTMATH)(_exp_xp)
+
+def exp(x):
+    """Compute the complex exponential of a real array."""
+    if backend_name() == "mlx.core":
+        return _exp_mx(x)
+    elif backend_name() == "numba":
+        return _exp_nb(x)
+    else:
+        return _exp_xp(x)
+