dspin 1.4.2__tar.gz → 1.5.0__tar.gz

{dspin-1.4.2/dspin.egg-info → dspin-1.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dspin
-Version: 1.4.2
+Version: 1.5.0
 Summary: Regulatory network models from single-cell perturbation profiling
 Home-page: https://github.com/JialongJiang/DSPIN
 Author: Jialong Jiang

{dspin-1.4.2 → dspin-1.5.0}/dspin/compute.py

@@ -21,6 +21,10 @@ from typing import Tuple, List, Callable, Any, Dict
 import warnings
 from scipy.sparse import issparse
 from joblib import Parallel, delayed, parallel_backend
+from sklearn.cluster import AgglomerativeClustering
+import leidenalg as la
+import igraph as ig
+import networkx as nx
 
 
 def category_balance_number(
@@ -70,11 +74,6 @@ def category_balance_number(
     return sampling_number
 
 
-from sklearn.cluster import AgglomerativeClustering
-import leidenalg as la
-import igraph as ig
-import networkx as nx
-
 def summary_components(all_components: np.array,
                        num_spin: int,
                        summary_method: str = 'kmeans',
@@ -125,8 +124,6 @@ def summary_components(all_components: np.array,
             # Determine gene cluster assignments based on the maximum component values.
             gene_groups_ind.append(np.argmax(components_kmeans, axis=0) == ii)
 
-        
-
         sc.set_figure_params(figsize=(8, 4))
 
         plt.subplot(1, 2, 1)
@@ -569,6 +566,7 @@ def para_moments_numpy(j_mat: np.array, h_vec: np.array) -> Tuple[np.array, np.a
 
     return corr_para, mean_para
 
+
 @numba.njit(fastmath=True, cache=True)
 def para_moments(j_mat: np.ndarray,
                        h_vec: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
@@ -708,8 +706,7 @@ def pseudol_gradient(cur_j: np.ndarray,
 def samp_moments(j_mat: np.ndarray, 
                  h_vec: np.ndarray, 
                  sample_size: int, 
-                 mixing_time: int, 
-                 samp_gap: int) -> Tuple[np.ndarray, np.ndarray]:
+                 mixing_time: int) -> Tuple[np.ndarray, np.ndarray]:
     """
     Sample moments for the Markov Chain Monte Carlo (MCMC).
 
@@ -723,8 +720,6 @@ def samp_moments(j_mat: np.ndarray,
         Size of the sample.
     mixing_time : int
         Mixing time for the MCMC.
-    samp_gap : int
-        Gap between samples.
 
     Returns
     -------
@@ -732,73 +727,241 @@ def samp_moments(j_mat: np.ndarray,
         The correlation parameter.
         The mean parameter.
     """
-    per_batch = int(1e5)
+    sample_size = int(sample_size)
+    mixing_time = int(mixing_time)
+
     num_spin = j_mat.shape[0]
-    rec_corr = np.zeros((num_spin, num_spin))
-    rec_mean = np.zeros(num_spin)
-    beta = 1
-    batch_count = 1
-    rec_sample = np.empty((num_spin, min(per_batch, int(sample_size))))
-    cur_spin = (np.random.randint(0, 3, (num_spin, 1)) - 1).astype(np.float64)
-    tot_sampling = int(
-        mixing_time + sample_size * samp_gap - mixing_time % samp_gap)
+    rec_corr = np.zeros((num_spin, num_spin), dtype=np.float64)
+    rec_mean = np.zeros(num_spin, dtype=np.float64)
+    beta = 1.0
 
-    rand_ind = np.random.randint(0, num_spin, tot_sampling)
-    rand_flip = np.random.randint(0, 2, tot_sampling)
+    # Cast & make contiguous (Numba-safe)
+    j_mat = np.ascontiguousarray(j_mat.astype(np.float64))
+    h_vec = np.ascontiguousarray(h_vec.astype(np.float64)).reshape(-1)
+    
+    # Spins in {-1,0,1}
+    cur_spin = np.sign(h_vec).astype(np.float64)
+
+    # Local field f = h + J @ s (update incrementally on changes)
+    f_vec = h_vec + j_mat.dot(cur_spin)
+
+    # Sampling schedule
+    tot_sampling = sample_size + mixing_time
+
+    # Pre-generate randomness (one uniform per site update)
     rand_prob = np.random.rand(tot_sampling)
 
-    # Monte Carlo Sampling
+    # Lazy accumulation buffers
+    last_outer = np.zeros((num_spin, num_spin), dtype=np.float64)  # current s s^T
+    have_last_outer = False
+    recording_started = False
+    weight = 0  # how many identical samples to add for current last_outer
+
+    # MCMC with sweep updates + heat-bath conditionals
     for ii in range(tot_sampling):
-        cur_ind = rand_ind[ii]
-        j_sub = j_mat[cur_ind, :]
-        accept_prob = 0.0
-        new_spin = 0.0
-        diff_energy = 0.0
-
-        if cur_spin[cur_ind] == 0:
-            if rand_flip[ii] == 0:
-                new_spin = 1.0
-            else:
-                new_spin = -1.0
-            diff_energy = -j_mat[cur_ind, cur_ind] - new_spin * (j_sub.dot(cur_spin) + h_vec[cur_ind])
-            accept_prob = min(1.0, np.exp(- diff_energy * beta)[0])
+        # sweep index (no random site index)
+        cur_ind = ii % num_spin
+
+        # current spin and local field excluding self-term
+        s_i = cur_spin[cur_ind]
+        Jii = j_mat[cur_ind, cur_ind]
+        hstar = f_vec[cur_ind] - Jii * s_i
+
+        # heat-bath weights for k in {+1, 0, -1}:
+        # log-weights: a = beta*( hstar + 0.5*Jii ), b = 0, c = beta*( -hstar + 0.5*Jii )
+        a = beta * (hstar + Jii)
+        b = 0.0
+        c = beta * (-hstar + Jii)
+        m = max(a, b, c)
+        pa = np.exp(a - m)
+        pb = np.exp(b - m)  # == exp(-m)
+        pc = np.exp(c - m)
+        Z = pa + pb + pc
+
+        # sample new spin
+        u = rand_prob[ii]
+        thresh_a = pa / Z
+        thresh_ab = (pa + pb) / Z
+        if u < thresh_a:
+            new_spin = 1.0
+        elif u < thresh_ab:
+            new_spin = 0.0
         else:
-            if rand_flip[ii] == 0:
-                accept_prob = 0;
-            else:
-                diff_energy = cur_spin[cur_ind] * (j_sub.dot(cur_spin) + h_vec[cur_ind])
-                accept_prob = min(1.0, np.exp(- diff_energy * beta)[0])
-
-        if rand_prob[ii] < accept_prob:
-            if cur_spin[cur_ind] == 0:
-                cur_spin[cur_ind] = new_spin
-            else:
-                cur_spin[cur_ind] = 0
-
-        if ii > mixing_time:
-            if (ii - mixing_time) % samp_gap == 0:
-                rec_sample[:, batch_count - 1] = cur_spin[:, 0].copy()
-                batch_count += 1
-
-                if batch_count == per_batch + 1:
-                    batch_count = 1
-                    rec_sample = np.ascontiguousarray(rec_sample)
-                    rec_corr += rec_sample.dot(rec_sample.T)
-                    rec_mean += np.sum(rec_sample, axis=1)
-
-    # Final processing of collected samples
-    if batch_count != 1:
-        cur_sample = rec_sample[:, :batch_count - 1]
-        cur_sample = np.ascontiguousarray(cur_sample)
-        rec_corr += cur_sample.dot(cur_sample.T)
-        rec_mean += np.sum(cur_sample, axis=1)
+            new_spin = - 1.0
+
+        # if spin changed, update cur_spin and cached field
+        changed = (new_spin != s_i)
+        if changed:
+            delta_s = new_spin - s_i
+            cur_spin[cur_ind] = new_spin
+            # incremental local-field update: f_k += J[k,i] * delta_s
+            for k in range(num_spin):
+                f_vec[k] += j_mat[k, cur_ind] * delta_s
+
+            # If we're recording and have accumulated weight, flush rec_corr BEFORE changing last_outer
+            if recording_started and weight > 0 and have_last_outer:
+                for k in range(num_spin):
+                    row = rec_corr[k]
+                    lo_row = last_outer[k]
+                    for l in range(num_spin):
+                        row[l] += weight * lo_row[l]
+                weight = 0
+
+            # Update last_outer incrementally if it exists
+            if recording_started and have_last_outer:
+                d = delta_s
+                i = cur_ind
+                # s's' = s s^T + d (e_i s'^T + s' e_i^T) - d^2 e_i e_i^T  (uses updated cur_spin)
+                for j in range(num_spin):
+                    last_outer[i, j] += d * cur_spin[j]
+                for j in range(num_spin):
+                    last_outer[j, i] += d * cur_spin[j]
+                last_outer[i, i] -= d * d
+
+        # Record sample on schedule
+        if ii >= mixing_time:
+            # mean: O(N)
+            for k in range(num_spin):
+                rec_mean[k] += cur_spin[k]
+
+            # initialize last_outer once at the first record
+            if not recording_started:
+                for k in range(num_spin):
+                    sk = cur_spin[k]
+                    row = last_outer[k]
+                    for l in range(num_spin):
+                        row[l] = sk * cur_spin[l]
+                have_last_outer = True
+                recording_started = True
+
+            # lazily count how many times to add current outer
+            weight += 1
+
+    # Flush any remaining weight after the loop
+    if recording_started and weight > 0 and have_last_outer:
+        for k in range(num_spin):
+            row = rec_corr[k]
+            lo_row = last_outer[k]
+            for l in range(num_spin):
+                row[l] += weight * lo_row[l]
 
     corr_para = rec_corr / sample_size
     mean_para = rec_mean / sample_size
-
+    
     return corr_para, mean_para
 
 
+@numba.njit()
+def samp_mcmc_states(j_mat: np.ndarray, 
+                    h_vec: np.ndarray, 
+                    sample_size: int, 
+                    mixing_time: int, 
+                    samp_gap: int) -> np.ndarray:
+    """
+    Sample states for the Markov Chain Monte Carlo (MCMC).
+
+    Parameters
+    ----------
+    j_mat : np.ndarray
+        Interaction matrix.
+    h_vec : np.ndarray
+        External field vector.
+    sample_size : int
+        Size of the sample.
+    mixing_time : int
+        Mixing time for the MCMC.
+    samp_gap : int
+        Gap between samples.
+
+    Returns
+    -------
+    np.ndarray
+        The sampled states.
+    """
+    sample_size = int(sample_size)
+    mixing_time = int(mixing_time)
+    samp_gap = int(samp_gap)
+
+    num_spin = j_mat.shape[0]
+
+    beta = 1.0
+    batch_count = 1
+    rec_sample = np.empty((num_spin, int(sample_size)), dtype=np.float64)
+
+    # Cast & make contiguous (Numba-safe)
+    j_mat = np.ascontiguousarray(j_mat.astype(np.float64))
+    # Accept both (n,) and (n,1) h_vec; flatten to 1D
+    h_vec = h_vec.reshape(h_vec.size)
+    h_vec = np.ascontiguousarray(h_vec.astype(np.float64))
+
+    # Spins in {-1,0,1}
+    cur_spin = (np.random.randint(0, 3, num_spin) - 1).astype(np.float64)
+
+    # Local field f = h + J @ s (update incrementally on changes)
+    f_vec = h_vec + j_mat.dot(cur_spin)
+
+    # Sampling schedule (avoid modulo inside loop)
+    rem = mixing_time % samp_gap
+    first_sample = mixing_time if rem == 0 else (mixing_time + (samp_gap - rem))
+    last_sample = first_sample + (sample_size - 1) * samp_gap
+    tot_sampling = last_sample + 1
+    next_sample = first_sample
+
+    # Pre-generate randomness: one uniform per single-site update
+    rand_prob = np.random.rand(tot_sampling)
+
+    # MCMC (Gibbs) with deterministic sweep order
+    for ii in range(tot_sampling):
+        cur_ind = ii % num_spin
+
+        # ---- Compute conditional for s_i given s_-i (PRE-UPDATE neighborhood) ----
+        s_i = cur_spin[cur_ind]
+        Jii = j_mat[cur_ind, cur_ind]
+        # h*_i = h_i + sum_{j!=i} J_ij s_j = f_i - J_ii * s_i  (pre-update)
+        hstar = f_vec[cur_ind] - Jii * s_i
+
+        # Unnormalized log-weights for k in {+1, 0, -1}
+        # Matches the enumerator convention (diagonal contributes +J_ii*k^2 here)
+        a = beta * (hstar + Jii)    # log w(+1)
+        b = 0.0                     # log w(0)
+        c = beta * (-hstar + Jii)   # log w(-1)
+
+        # Stable normalization via max-trick
+        m = max(a, b, c)
+        pa = np.exp(a - m)
+        pb = np.exp(b - m)
+        pc = np.exp(c - m)
+        Z = pa + pb + pc
+
+        # Sample new spin from the conditional (heat-bath)
+        u = rand_prob[ii]
+        t1 = pa / Z
+        t2 = (pa + pb) / Z
+        if u < t1:
+            new_spin = 1.0
+        elif u < t2:
+            new_spin = 0.0
+        else:
+            new_spin = -1.0
+
+        # If changed, update spin and cached local field
+        if new_spin != s_i:
+            delta_s = new_spin - s_i
+            cur_spin[cur_ind] = new_spin
+            # Incremental local-field update: f_k += J[k, i] * delta_s
+            for k in range(num_spin):
+                f_vec[k] += j_mat[k, cur_ind] * delta_s
+
+        # Record sample on schedule
+        if ii == next_sample:
+            for k in range(num_spin):
+                rec_sample[k, batch_count - 1] = cur_spin[k]
+            batch_count += 1
+            next_sample += samp_gap
+
+    return rec_sample
+
+
 def _gradient_one_round(kk: int,
                         cur_j: np.ndarray,
                         cur_h: np.ndarray,
@@ -814,7 +977,7 @@ def _gradient_one_round(kk: int,
             corr_para, mean_para = para_moments(cur_j, cur_h[:, kk])
         elif method == 'mcmc_maximum_likelihood':
             corr_para, mean_para = samp_moments(
-                    cur_j, cur_h[:, kk], train_dat['mcmc_samplingsz'], train_dat['mcmc_samplingmix'], train_dat['mcmc_samplegap'])
+                    cur_j, cur_h[:, kk], train_dat['mcmc_samplingsz'], train_dat['mcmc_samplingmix'])
         j_grad  = corr_para - raw_data_k[0]
         h_grad  = mean_para  - raw_data_k[1].flatten()
 
@@ -985,6 +1148,7 @@ def apply_regularization(rec_jgrad: np.ndarray,
 
         if 'lambda_l2_h_rela_prior' in train_dat:
             perturb_matrix = train_dat['perturb_matrix']
+            rec_hgrad = average_batch_control(rec_hgrad, if_control, batch_index)
             rec_hgrad += train_dat['lambda_l2_h_rela_prior'] * (h_rela - perturb_matrix)
 
         if 'lambda_l1_h_rela' in train_dat:
@@ -1421,7 +1585,7 @@ def compute_relative_responses(cur_h: np.ndarray,
         batch_samples_idx = np.where(batch_index == current_batch)[0]
 
         if np.any(if_control[batch_samples_idx]):
-            batch_controls_idx = batch_samples_idx[if_control[batch_samples_idx]]
+            batch_controls_idx = batch_samples_idx[np.where(if_control[batch_samples_idx])[0]]
             batch_controls_mean = np.mean(cur_h[:, batch_controls_idx], axis=1)
         else:
             batch_controls_mean = all_controls_mean
@@ -1432,6 +1596,40 @@ def compute_relative_responses(cur_h: np.ndarray,
     return relative_h
 
 
+def average_batch_control(hgrad_raw: np.ndarray, 
+                        if_control: np.ndarray, 
+                        batch_index: np.ndarray) -> np.ndarray:
+    """
+    Assign the gradient of control samples as the average gradient of the samples in the same batch.
+
+    Parameters
+    ----------
+    hgrad : np.ndarray
+        The gradient of the h matrix.
+    if_control : np.ndarray
+        A boolean array indicating which samples are control samples.
+    batch_index : np.ndarray
+        An array indicating the batch assignment for each sample.
+    """
+
+    hgrad = hgrad_raw.copy() 
+    
+    unique_batches = np.unique(batch_index)
+    num_batches = len(unique_batches)
+
+    for current_batch in unique_batches:
+        batch_samples_idx = np.where(batch_index == current_batch)[0]
+        
+        if np.any(if_control[batch_samples_idx]):
+            if np.any(if_control[batch_samples_idx]):
+                batch_controls_idx = batch_samples_idx[np.where(if_control[batch_samples_idx])[0]]
+                batch_mean = np.mean(hgrad_raw[:, batch_samples_idx], axis=1)
+
+                hgrad[:, batch_controls_idx] = batch_mean[:, np.newaxis]
+
+    return hgrad
+
+
 def select_representative_sample(raw_data: List[Tuple[np.ndarray, np.ndarray]], 
                                  num_select: int) -> List[int]:
     """

{dspin-1.4.2 → dspin-1.5.0/dspin.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dspin
-Version: 1.4.2
+Version: 1.5.0
 Summary: Regulatory network models from single-cell perturbation profiling
 Home-page: https://github.com/JialongJiang/DSPIN
 Author: Jialong Jiang

{dspin-1.4.2 → dspin-1.5.0}/setup.py

@@ -21,7 +21,7 @@ install_requires = [
 # ]
 
 # Managing your version (adjust accordingly)
-version = "1.4.2"  # You can manage versioning in a more sophisticated way if needed
+version = "1.5.0"  # You can manage versioning in a more sophisticated way if needed
 
 # Long description (usually a README file)
 readme = open("README_PyPI.md").read()  # Adjust the file name if necessary