vbi 0.1.3__cp310-cp310-manylinux2014_x86_64.whl → 0.2__cp310-cp310-manylinux2014_x86_64.whl

vbi/models/cupy/km.py

@@ -4,40 +4,61 @@ from vbi.models.cupy.utils import *
 
 
 class KM_sde:
-
-    valid_parameters = [
-        "num_sim",        # number of simulations
-        "G",              # global coupling strength
-        "dt",             # time step
-        "noise_amp",      # noise amplitude
-        "omega",          # natural angular frequency
-        "weights",        # weighted connection matrix
-        "seed",
-        "alpha",          # frustration matrix
-        "t_initial",      # initial time
-        "t_transition",   # transition time
-        "t_end",          # end time
-        "output",         # output directory
-        "num_threads",    # number of threads using openmp
-        "initial_state",
-        "type",           # output times series data type
-        "engine",         # cpu or gpu
-    ]
+    
+    """
+    Kuramoto model with noise (stochastic differential equation)
+    
+    Parameters
+    ----------
+    G: float
+        global coupling strength
+    dt: float
+        time step
+    noise_amp: float
+        noise amplitude
+    weights: array
+        weighted connection matrix
+    omega: array
+        natural angular frequency
+    seed: int
+        fix random seed for initial state
+    t_cut: float
+        transition time
+    t_end: float
+        end time
+    decimate: int
+        decimate the output time series
+    output: str
+        output directory
+    initial_state: array
+        initial state
+    engine: str
+        cpu or gpu
+    type: type
+        data type for calculations, default is np.float32
+    alpha: array # TODO not implemented
+        frustration matrix 
+    num_sim: int
+        number of simulations
+    method: str
+        integration method, default is heun
+    same_initial_state: bool
+        use the same initial state for all simulations, default is False    
+    
+    """
 
     def __init__(self, par={}) -> None:
 
-        self.check_parameters(par)
         self._par = self.get_default_parameters()
+        self.valid_parameters = list(self._par.keys())
+        self.check_parameters(par)
         self._par.update(par)
 
         for item in self._par.items():
-            name = item[0]
-            value = item[1]
-            setattr(self, name, value)
+            setattr(self, item[0], item[1])
 
         self.xp = get_module(self.engine)
         self.ns = self.num_sim
-
         self.nn = self.num_nodes = self.weights.shape[0]
 
         if self.seed is not None:
@@ -55,9 +76,7 @@ class KM_sde:
         print (f"Kuramoto model with noise (sde), {self.engine} implementation.")
         print ("----------------")
         for item in self._par.items():
-            name = item[0]
-            value = item[1]
-            print (f"{name} = {value}")
+            print (f"{item[0]} = {item[1]}")
         return ""
 
     def __call__(self):
@@ -74,18 +93,18 @@ class KM_sde:
             "weights": None,                 # weighted connection matrix
             "omega": None,                   # natural angular frequency
             "seed": None,                    # fix random seed for initial state
-            "t_initial": 0.0,                # initial time
-            "t_transition": 0.0,             # transition time
+            "t_start": 0.0,                  # initial time
+            "t_cut": 0.0,                    # transition time
             "t_end": 100.0,                  # end time
-            "num_threads": 1,                # number of threads using openmp
             "output": "output",              # output directory
             "initial_state": None,           # initial state
             "engine": "cpu",                 # cpu or gpu
             "type": np.float32,              # output times series data type
-            "alpha": None,                   # frustration matrix
+            "alpha": None,                   # frustration matrix # TODO not implemented
             "num_sim": 1,                    # number of simulations
             "method": "heun",                # integration method
             "same_initial_state": False,     # use the same initial state for all simulations
+            "decimate": 1,                   # decimate the output time series
 
         }
 
@@ -130,12 +149,12 @@ class KM_sde:
         x = self.initial_state
         xs = []
         integrator = self.euler if self.method == "euler" else self.heun
-        n_transition = int(self.t_transition /
-                           self.dt) if self.t_transition > 0 else 1
+        n_transition = int(self.t_cut /
+                           self.dt) if self.t_cut > 0 else 1
 
         for it in tqdm.tqdm(range(1, len(t)), disable=not verbose, desc="Integrating"):
             x = integrator(x, t[it])
-            if it >= n_transition:
+            if (it >= n_transition) and (it % self.decimate == 0):
                 if self.engine == "gpu":
                     xs.append(x.get())
                 else:
@@ -144,6 +163,15 @@ class KM_sde:
         t = t[n_transition:]
 
         return {"t": t, "x": xs}
+    
+    def step(self, x, t):
+        ''' Step function for the model'''
+        if self.method == "euler":
+            return self.euler(x, t)
+        elif self.method == "heun":
+            return self.heun(x, t)
+        else:
+            raise ValueError(f"Invalid method: {self.method}")
 
     def run(self, x0=None, verbose=True):
         '''

vbi/models/cupy/mpr.py

@@ -150,7 +150,7 @@ class MPR_sde:
         model parameter
     tr: float
         repetition time of fMRI
-    noise_amp: float
+    noise_amp: float, np.array
         amplitude of noise
     same_noise_per_sim:
         same noise for all simulations
@@ -200,6 +200,8 @@ class MPR_sde:
             name = item[0]
             value = item[1]
             setattr(self, name, value)
+        
+        self.update_dependent_parameters()
             
         self.B = Bold(Bpar)
 
@@ -266,7 +268,7 @@ class MPR_sde:
             "same_initial_state": False,
         }
         dt = par["dt"]
-        noise_amp = par["noise_amp"]
+        noise_amp = np.array(par["noise_amp"])
         sigma_r = np.sqrt(dt) * np.sqrt(2 * noise_amp)
         sigma_v = np.sqrt(dt) * np.sqrt(4 * noise_amp)
         par["sigma_r"] = sigma_r
@@ -274,17 +276,35 @@ class MPR_sde:
         # par.update(self.get_balloon_parameters())
 
         return par
+    
+    def update_dependent_parameters(self):
+        dt = self.dt
+        noise_amp = self.noise_amp
+        if hasattr(noise_amp, "__iter__"):
+            noise_amp = np.array(noise_amp)
+        else:
+            noise_amp = np.array([noise_amp])
+            
+        sigma_r = np.sqrt(dt) * np.sqrt(2 * noise_amp)
+        sigma_v = np.sqrt(dt) * np.sqrt(4 * noise_amp)
+        self.sigma_r = sigma_r
+        self.sigma_v = sigma_v
+        self._par["sigma_r"] = sigma_r
+        self._par["sigma_v"] = sigma_v
+
 
     def check_parameters(self, par):
         for key in par.keys():
             if key not in self.valid_parameters:
                 raise ValueError(f"Invalid parameter {key:s} provided.")
-
+        
     def prepare_input(self):
 
         self.G = self.xp.array(self.G)
+        self.sigma_r = self.xp.array(self.sigma_r)
+        self.sigma_v = self.xp.array(self.sigma_v)
         assert self.weights is not None, "weights must be provided"
-        self.weights = self.xp.array(self.weights).T  # ! Directed network #!TODO: check
+        self.weights = self.xp.array(self.weights)
         self.weights = move_data(self.weights, self.engine)
         self.nn = self.num_nodes = self.weights.shape[0]
 

vbi/models/cupy/utils.py

@@ -75,6 +75,47 @@ def todevice(x):
     '''
     return cp.asarray(x)
 
+# write a function to detexct where is x on cpu or gpu
+
+def is_on_cpu(x):
+    '''
+    Check if input is on CPU (i.e., not a CuPy array)
+
+    Parameters
+    ----------
+    x: any
+        Input to check
+
+    Returns
+    -------
+    bool
+        True if input is not a CuPy array (i.e., is on CPU), False otherwise
+    '''
+    if cp is None:
+        return None
+    return not isinstance(x, cp.ndarray)
+
+def is_on_gpu(x):
+    '''
+    Check if input is on GPU (i.e., is a CuPy array)
+
+    Parameters
+    ----------
+    x: any
+        Input to check
+
+    Returns
+    -------
+    bool or None
+        True if input is a CuPy array (i.e., is on GPU)
+        False if input is not a CuPy array
+        None if CuPy is not installed
+    '''
+    if cp is None:
+        return None
+    return isinstance(x, cp.ndarray)
+
+
 
 def move_data(x, engine):
     if engine == "cpu":
@@ -124,6 +165,10 @@ def is_seq(x):
     '''
     return hasattr(x, '__iter__')
 
+# def is_seq(x):
+#     """Check if x is a sequence (list, tuple, array) but not a string"""
+#     return hasattr(x, '__len__') and not isinstance(x, (str, bytes))
+
 
 def prepare_vec(x, ns, engine, dtype="float"):
     '''
@@ -137,6 +182,8 @@ def prepare_vec(x, ns, engine, dtype="float"):
         number of simulations
     engine: str
         cpu or gpu
+    dtype: str or numpy.dtype
+        data type to convert to
 
     Returns
     -------
@@ -145,9 +192,18 @@ def prepare_vec(x, ns, engine, dtype="float"):
 
     '''
     xp = get_module(engine)
-
     if not is_seq(x):
-        return eval(f"{dtype}({x})")
+        # Convert dtype string to numpy dtype
+        if isinstance(dtype, str):
+            if dtype == "float":
+                dtype = np.float64
+            elif dtype == "float32":
+                dtype = np.float32
+            elif dtype == "float64":
+                dtype = np.float64
+            else:
+                raise ValueError(f"Unsupported dtype: {dtype}")
+        x = xp.array(x, dtype=dtype)
     else:
         x = np.array(x)
         if x.ndim == 1:
@@ -159,7 +215,31 @@ def prepare_vec(x, ns, engine, dtype="float"):
             raise ValueError("x.ndim must be 1 or 2")
     return x.astype(dtype)
 
+def prepare_vec_1d(x, ns, engine, dtype="float"):
+    '''
+    check and prepare vector dimension and type
+    '''
+    if not is_seq(x):
+        # Convert dtype string to numpy dtype
+        if isinstance(dtype, str):
+            if dtype == "float":
+                dtype = np.float64
+            elif dtype == "float32":
+                dtype = np.float32
+            elif dtype == "float64":
+                dtype = np.float64
+            else:
+                raise ValueError(f"Unsupported dtype: {dtype}")
+        x = np.array(x, dtype=dtype)
+    else:
+        assert(x.ndim == 1), "x must be a 1d array"
+        assert(x.shape[0] == ns), "x must have the same number of elements as ns"
+    x = move_data(x, engine)
+    return x.astype(dtype)
 
+                    
+        
+    
 def get_(x, engine="cpu", dtype="f"):
     """
     Parameters
@@ -182,3 +262,84 @@ def get_(x, engine="cpu", dtype="f"):
         return x.get().astype(dtype)
     else:
         return x.astype(dtype)
+
+
+def dtype_convert(dtype):
+    """
+    Convert a string representation of a data type to a numpy dtype.
+
+    Parameters
+    ----------
+    dtype : str
+        The string representation of the data type (e.g., "float", "float32", "float64").
+
+    Returns
+    -------
+    numpy.dtype
+        The corresponding numpy dtype.
+    
+    Raises
+    ------
+    ValueError
+        If the input string does not match any known data type.
+    """
+    
+    if dtype == "float":
+        return np.float64
+    elif dtype == "f":
+        return np.float32
+    elif dtype == "float32":
+        return np.float32
+    elif dtype == "float64":
+        return np.float64
+    else:
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+
+def prepare_vec_2d(x, nn, ns, engine, dtype="float"):
+    ''' 
+    if x is scalar pass 
+    if x is 1d array, shape should be (ns,)
+    if x is 2d array, shape should be (nn, ns)
+    '''
+    import numpy as np
+    
+    # Determine target numpy dtype
+    try:
+        target_dtype = dtype_convert(dtype)
+    except ValueError:
+        target_dtype = np.dtype(dtype)
+    
+    if not is_seq(x):
+        # x is scalar - return
+        return x
+    
+    # Convert x to numpy array if it isn't already
+    x = np.asarray(x)
+    
+    if x.ndim == 1:
+        # 1D array - should have shape (ns,)
+        if x.shape[0] != ns:
+            raise ValueError(f"1D array should have shape ({ns},), got {x.shape}")
+        
+        # Change dtype
+        x = x.astype(target_dtype)
+        
+        # Move to appropriate device
+        x = move_data(x, engine)
+        
+    elif x.ndim == 2:
+        # 2D array - should have shape (nn, ns)
+        if x.shape != (nn, ns):
+            raise ValueError(f"2D array should have shape ({nn}, {ns}), got {x.shape}")
+        
+        # Change dtype
+        x = x.astype(target_dtype)
+        
+        # Move to appropriate device
+        x = move_data(x, engine)
+        
+    else:
+        raise ValueError(f"Array should be 1D or 2D, got {x.ndim}D")
+    
+    return x